Several improvements to capabilities and build.
- cmake fixed, no longer needs numpy headers. - models can be loaded from an external interepreter.
This commit is contained in:
Nicholas Lubbers
@ -1,12 +1,3 @@
|
||||
if(CMAKE_VERSION VERSION_LESS 3.12)
|
||||
#This block was not tested, mimmicks PYTHON.cmake.
|
||||
find_package(PythonLibs REQUIRED) # Deprecated since version 3.12
|
||||
target_include_directories(lammps PRIVATE ${PYTHON_INCLUDE_DIR})
|
||||
target_link_libraries(lammps PRIVATE ${PYTHON_LIBRARY})
|
||||
target_include_directories(lammps PRIVATE ${Python_NumPy_INCLUDE_DIRS})
|
||||
else()
|
||||
find_package(Python REQUIRED COMPONENTS NumPy)
|
||||
target_include_directories(lammps PRIVATE ${Python_NumPy_INCLUDE_DIRS})
|
||||
endif()
|
||||
target_compile_definitions(lammps PRIVATE -DLMP_MLIAPPY)
|
||||
execute_process(COMMAND cythonize mliap_model_python_couple.pyx WORKING_DIRECTORY ../src/MLIAPPY)
|
||||
|
||||
target_compile_definitions(lammps PRIVATE -DLMP_MLIAPPY)
|
||||
|
||||
@ -3,7 +3,8 @@ if(CMAKE_VERSION VERSION_LESS 3.12)
|
||||
target_include_directories(lammps PRIVATE ${PYTHON_INCLUDE_DIRS})
|
||||
target_link_libraries(lammps PRIVATE ${PYTHON_LIBRARIES})
|
||||
else()
|
||||
find_package(Python REQUIRED COMPONENTS Development)
|
||||
find_package(Python REQUIRED COMPONENTS Development Interpreter)
|
||||
target_include_directories(lammps PRIVATE ${Python_INCLUDE_DIRS})
|
||||
target_link_libraries(lammps PRIVATE Python::Python)
|
||||
endif()
|
||||
target_compile_definitions(lammps PRIVATE -DLMP_PYTHON)
|
||||
|
||||
@ -692,12 +692,10 @@ Extension to the MLIAP package for coupling with python models.
|
||||
|
||||
To use this package, also the :ref:`MLIAP package <PKG-MLIAP>` needs to be installed.
|
||||
To use this package, also the :ref:`PYTHON package <PKG-PYTHON>` needs to be installed.
|
||||
The version of python must be >3.5, and has been tested only with 3.8.
|
||||
Compiling this package has only been tested using CMake, not with pure makefiles.s
|
||||
The python interpreter linked to LAMMPS will need cython and numpy installed.
|
||||
The version of python must be >3.5.
|
||||
|
||||
Before compiling, run cythonize on /src/MLIAPPY/mliap_model_python_couple.pyx.
|
||||
This will produce /src/MLIAPPY/mliap_model_python_couple.cpp and /src/MLIAPPY/mliap_model_python_couple.h files.
|
||||
The python interpreter linked to LAMMPS will need cython and numpy installed.
|
||||
The examples build models with pytorch, which would thus need to be installed.
|
||||
|
||||
This package includes more options for the mliap compute and pair style.
|
||||
|
||||
@ -706,6 +704,7 @@ This package includes more options for the mliap compute and pair style.
|
||||
**Supporting info:**
|
||||
|
||||
* src/MLIAPPY: filenames -> commands
|
||||
* src/MLIAPPY/README
|
||||
* :doc:`pair_style mliap <pair_mliap>`
|
||||
* examples/mliappy (see README)
|
||||
----------
|
||||
|
||||
26
examples/mliappy/README
Normal file
26
examples/mliappy/README
Normal file
@ -0,0 +1,26 @@
|
||||
README for MLIAPPY Example
|
||||
|
||||
These examples run the Ta06 example from the MLIAP package, but using the python coupling.
|
||||
|
||||
1: Running models using LAMMPS executable: in.mliap.snap.Ta06A.
|
||||
|
||||
To run this, first run convert_mliap_Ta06A.py, which will convert the Ta06 potential into a pytorch model.
|
||||
It will be saved as "Ta06A.mliap.pytorch.model.pkl".
|
||||
|
||||
It will also copy the "mliappy_pytorch.py" file into the current working directory. mliappy_pytorch.py contains
|
||||
class definitions suitable for wrapping an arbitrary energy model MLIAPPY. It must be available to python when
|
||||
creating or unpicking a pytorch model for MLIAPPY.
|
||||
|
||||
From that point you can run the example lmp -in in.mliap.snap.Ta06A -echo both
|
||||
|
||||
2: Running models from python with LAMMPS in library mode: load_external.py
|
||||
|
||||
Before testing this, ensure that the first example (using LAMMPS executable) works.
|
||||
|
||||
Not all python installations support this mode of operation.
|
||||
Too test if your interpreter supports this, run:
|
||||
`python test_pylibs.py`
|
||||
and examine the output.
|
||||
|
||||
If this succeeds, you should be able to run:
|
||||
`python load_external.py`
|
||||
@ -1,11 +0,0 @@
|
||||
README for MLIAPPY Example
|
||||
|
||||
This example runs the Ta06 example from the MLIAP example, but using the python coupling.
|
||||
|
||||
To run this, first run convert_mliap_Ta06A.py, which will convert the Ta06 potential into a pytorch model.
|
||||
It will be saved as "Ta06A.mliap.pytorch.model.pkl".
|
||||
|
||||
It will also copy the "torchlink.py" file into the current working directory. torchlink.py contains
|
||||
class definitions suitable for wrapping an arbitrary energy model MLIAPPY.
|
||||
|
||||
From that point you can run the example lmp -in in.mliap.snap.Ta06A -echo both
|
||||
@ -22,7 +22,7 @@ with torch.autograd.no_grad():
|
||||
lin.bias.set_(torch.as_tensor(bias,dtype=torch.float64).unsqueeze(0))
|
||||
|
||||
# Wrap the pytorch model for usage with MLIAPPY
|
||||
model = mliappy_pytorch.IgnoreTypes(lin)
|
||||
model = mliappy_pytorch.IgnoreElems(lin)
|
||||
n_descriptors = lin.weight.shape[1]
|
||||
n_params = mliappy_pytorch.calc_n_params(model)
|
||||
n_types = 1
|
||||
@ -30,4 +30,4 @@ linked_model = mliappy_pytorch.TorchWrapper64(model,n_descriptors=n_descriptors,
|
||||
|
||||
# Save the result.
|
||||
with open("Ta06A.mliap.pytorch.model.pkl",'wb') as pfile:
|
||||
pickle.dump(linked_model,pfile)
|
||||
pickle.dump(linked_model,pfile)
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
# Demonstrate MLIAP interface to kinear SNAP potential
|
||||
# Demonstrate MLIAP interface to linear SNAP potential
|
||||
|
||||
# Initialize simulation
|
||||
|
||||
|
||||
104
examples/mliappy/load_external.py
Normal file
104
examples/mliappy/load_external.py
Normal file
@ -0,0 +1,104 @@
|
||||
|
||||
# Demonstrate how to load a model from the python side.
|
||||
# This is essentially the same as in.mliap.pytorch.Ta06A--
|
||||
# except that python is the driving program, and lammps
|
||||
# is in library mode.
|
||||
|
||||
before_loading =\
|
||||
"""
|
||||
# Demonstrate MLIAP interface to linear SNAP potential
|
||||
|
||||
# Initialize simulation
|
||||
|
||||
variable nsteps index 100
|
||||
variable nrep equal 4
|
||||
variable a equal 3.316
|
||||
units metal
|
||||
|
||||
# generate the box and atom positions using a BCC lattice
|
||||
|
||||
variable nx equal ${nrep}
|
||||
variable ny equal ${nrep}
|
||||
variable nz equal ${nrep}
|
||||
|
||||
boundary p p p
|
||||
|
||||
lattice bcc $a
|
||||
region box block 0 ${nx} 0 ${ny} 0 ${nz}
|
||||
create_box 1 box
|
||||
create_atoms 1 box
|
||||
|
||||
mass 1 180.88
|
||||
|
||||
# choose potential
|
||||
# DATE: 2014-09-05 UNITS: metal CONTRIBUTOR: Aidan Thompson athomps@sandia.gov CITATION: Thompson, Swiler, Trott, Foiles and Tucker, arxiv.org, 1409.3880 (2014)
|
||||
|
||||
# Definition of SNAP potential Ta_Cand06A
|
||||
# Assumes 1 LAMMPS atom type
|
||||
|
||||
variable zblcutinner equal 4
|
||||
variable zblcutouter equal 4.8
|
||||
variable zblz equal 73
|
||||
|
||||
# Specify hybrid with SNAP, ZBL
|
||||
|
||||
pair_style hybrid/overlay &
|
||||
zbl ${zblcutinner} ${zblcutouter} &
|
||||
mliap model mliappy LATER &
|
||||
descriptor sna Ta06A.mliap.descriptor
|
||||
pair_coeff 1 1 zbl ${zblz} ${zblz}
|
||||
pair_coeff * * mliap Ta
|
||||
"""
|
||||
after_loading =\
|
||||
"""
|
||||
|
||||
# Setup output
|
||||
|
||||
compute eatom all pe/atom
|
||||
compute energy all reduce sum c_eatom
|
||||
|
||||
compute satom all stress/atom NULL
|
||||
compute str all reduce sum c_satom[1] c_satom[2] c_satom[3]
|
||||
variable press equal (c_str[1]+c_str[2]+c_str[3])/(3*vol)
|
||||
|
||||
thermo_style custom step temp epair c_energy etotal press v_press
|
||||
thermo 10
|
||||
thermo_modify norm yes
|
||||
|
||||
# Set up NVE run
|
||||
|
||||
timestep 0.5e-3
|
||||
neighbor 1.0 bin
|
||||
neigh_modify once no every 1 delay 0 check yes
|
||||
|
||||
# Run MD
|
||||
|
||||
velocity all create 300.0 4928459 loop geom
|
||||
fix 1 all nve
|
||||
run ${nsteps}
|
||||
"""
|
||||
|
||||
|
||||
import lammps
|
||||
|
||||
lmp = lammps.lammps(cmdargs=['-echo','both'])
|
||||
# This commmand must be run before the MLIAP object is declared in lammps.
|
||||
lmp.mliappy.activate()
|
||||
|
||||
lmp.commands_string(before_loading)
|
||||
|
||||
# Now the model is declared, but empty -- because the model filename
|
||||
# was given as "LATER".
|
||||
|
||||
# Load the python module, construct on the fly, do whatever, here:
|
||||
import pickle
|
||||
with open('Ta06A.mliap.pytorch.model.pkl','rb') as pfile:
|
||||
model = pickle.load(pfile)
|
||||
|
||||
# Now that you have a model, connect it to the pairstyle
|
||||
lmp.mliappy.load_model(model)
|
||||
|
||||
# Proceed with whatever calculations you like.
|
||||
lmp.commands_string(after_loading)
|
||||
|
||||
|
||||
12
examples/mliappy/test_pylibs.py
Normal file
12
examples/mliappy/test_pylibs.py
Normal file
@ -0,0 +1,12 @@
|
||||
import sysconfig, os,ctypes
|
||||
|
||||
library = sysconfig.get_config_vars('INSTSONAME')[0]
|
||||
|
||||
pylib = ctypes.CDLL(library)
|
||||
|
||||
connected = pylib.Py_IsInitialized()
|
||||
|
||||
if not connected:
|
||||
print("FAILURE: This interpreter is not compatible with python-driven mliappy.")
|
||||
else:
|
||||
print("SUCCESS: This interpreter is compatible with python-driven MLIAPPY")
|
||||
@ -505,6 +505,8 @@ class lammps(object):
|
||||
self.FIX_EXTERNAL_CALLBACK_FUNC = CFUNCTYPE(None, py_object, self.c_bigint, c_int, POINTER(self.c_tagint), POINTER(POINTER(c_double)), POINTER(POINTER(c_double)))
|
||||
self.lib.lammps_set_fix_external_callback.argtypes = [c_void_p, c_char_p, self.FIX_EXTERNAL_CALLBACK_FUNC, py_object]
|
||||
self.lib.lammps_set_fix_external_callback.restype = None
|
||||
|
||||
self.mliappy = MLIAPPY(self)
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# shut-down LAMMPS instance
|
||||
@ -2924,3 +2926,43 @@ class IPyLammps(PyLammps):
|
||||
"""
|
||||
from IPython.display import HTML
|
||||
return HTML("<video controls><source src=\"" + filename + "\"></video>")
|
||||
|
||||
class MLIAPPY():
|
||||
def __init__(self,lammps):
|
||||
self._module = None
|
||||
self.lammps = lammps
|
||||
|
||||
@property
|
||||
def module(self):
|
||||
if self._module:
|
||||
return self._module
|
||||
|
||||
try:
|
||||
# Begin Importlib magic to find the embedded python module
|
||||
# This is needed because the filename for liblammps does not
|
||||
# match the spec for normal python modules, wherein
|
||||
# file names match with PyInit function names.
|
||||
# Also, python normally doesn't look for extensions besides '.so'
|
||||
# We fix both of these problems by providing an explict
|
||||
# path to the extension module 'mliap_model_python_couple' in
|
||||
import sys
|
||||
import importlib.util
|
||||
import importlib.machinery
|
||||
|
||||
path = self.lammps.lib._name
|
||||
loader = importlib.machinery.ExtensionFileLoader('mliap_model_python_couple',path)
|
||||
spec = importlib.util.spec_from_loader('mliap_model_python_couple',loader)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules['mliap_model_python_couple']=module
|
||||
spec.loader.exec_module(module)
|
||||
self._module = module
|
||||
# End Importlib magic to find the embedded python module
|
||||
except:
|
||||
raise ImportError("Could not load MLIAPPY coupling module")
|
||||
|
||||
def activate(self):
|
||||
self.module
|
||||
|
||||
def load_model(self,model):
|
||||
self.module.load_from_python(model)
|
||||
|
||||
|
||||
@ -69,6 +69,17 @@ PairMLIAP::~PairMLIAP()
|
||||
|
||||
void PairMLIAP::compute(int eflag, int vflag)
|
||||
{
|
||||
|
||||
// consistency checks
|
||||
|
||||
if (data->ndescriptors != model->ndescriptors) {
|
||||
error->all(FLERR,"Incompatible model and descriptor descriptor count");
|
||||
};
|
||||
|
||||
if (data->nelements != model->nelements) {
|
||||
error->all(FLERR,"Incompatible model and descriptor element count");
|
||||
};
|
||||
|
||||
ev_init(eflag,vflag);
|
||||
|
||||
data->generate_neighdata(list, eflag, vflag);
|
||||
@ -137,9 +148,8 @@ void PairMLIAP::settings(int narg, char ** arg)
|
||||
if (iarg+3 > narg) error->all(FLERR,"Illegal pair_style mliap command");
|
||||
model = new MLIAPModelQuadratic(lmp,arg[iarg+2]);
|
||||
iarg += 3;
|
||||
}
|
||||
#ifdef LMP_MLIAPPY
|
||||
else if (strcmp(arg[iarg+1],"mliappy") == 0) {
|
||||
} else if (strcmp(arg[iarg+1],"mliappy") == 0) {
|
||||
if (iarg+3 > narg) error->all(FLERR,"Illegal pair_style mliap command");
|
||||
model = new MLIAPModelPython(lmp,arg[iarg+2]);
|
||||
iarg += 3;
|
||||
@ -225,15 +235,7 @@ void PairMLIAP::coeff(int narg, char **arg)
|
||||
data = new MLIAPData(lmp, gradgradflag, map, model, descriptor, this);
|
||||
data->init();
|
||||
|
||||
// consistency checks
|
||||
|
||||
if (data->ndescriptors != model->ndescriptors) {
|
||||
error->all(FLERR,"Incompatible model and descriptor definitions (different number of descriptors)");
|
||||
};
|
||||
|
||||
if (data->nelements != model->nelements) {
|
||||
error->all(FLERR,"Incompatible model and descriptor definitions (different number of elements)");
|
||||
};
|
||||
}
|
||||
|
||||
/* ----------------------------------------------------------------------
|
||||
|
||||
1
src/MLIAPPY/Makefile.lammps
Normal file
1
src/MLIAPPY/Makefile.lammps
Normal file
@ -0,0 +1 @@
|
||||
#TODO
|
||||
37
src/MLIAPPY/README
Normal file
37
src/MLIAPPY/README
Normal file
@ -0,0 +1,37 @@
|
||||
README for MLIAPPY source files.
|
||||
|
||||
MLIAPPY requires python 3 with cython, numpy, and optionally pytorch installed.
|
||||
One could build compatible python models without pytorch, with a bit of work by hand.
|
||||
You can get these via standard procedures (pip, conda, or similar)
|
||||
|
||||
MLIAPPY also requires the LAMMPS packages MLIAP and PYTHON be enabled.
|
||||
|
||||
MLIAPPY can be built with cmake -- a pure make version is forthcoming.
|
||||
|
||||
If you are building LAMMPS manually (no cmake/make), the process works roughly as possible:
|
||||
|
||||
First run cythonize on mliap_model_python_couple.pyx. This will generate:
|
||||
|
||||
(1) mliap_model_python_couple.h
|
||||
(2) mliap_model_python_couple.cpp
|
||||
|
||||
File (1) is a roughly human readable header file, (2) is a large source file,
|
||||
and navigating it is not for the faint of heart. These files are used in mliap_model_python.cpp.
|
||||
|
||||
Then use compiler options to define the macro LMP_MLIAPPY during compilation.
|
||||
|
||||
Other information:
|
||||
|
||||
The "mliap_model_python.cpp" files and "mliap_model_python.h" files cover the
|
||||
definitions of the LAMMPS object MLIAPModelPython.
|
||||
|
||||
How does this all work?
|
||||
|
||||
Roughly: A python extension module called "mliap_model_python_couple" is built in to LAMMPS. This holds a dictionary of currently defined MLIAPModelPython objects and
|
||||
the equivalent python models. It also converts the data needed by the model from
|
||||
C arrays to numpy arrays before passing them to the python model.
|
||||
|
||||
The last file is mliappy_pytorch.py. This contains some simple utility classes
|
||||
for taking an energy model in pytorch, and wrapping it for compatibily with MLIAP.
|
||||
|
||||
For more information on MLIAPPY, see the examples in examples/mliappy.
|
||||
@ -29,28 +29,25 @@ using namespace LAMMPS_NS;
|
||||
MLIAPModelPython::MLIAPModelPython(LAMMPS* lmp, char* coefffilename) :
|
||||
MLIAPModel(lmp, coefffilename)
|
||||
{
|
||||
int err;
|
||||
|
||||
model_loaded = 0;
|
||||
python->init();
|
||||
|
||||
PyGILState_STATE gstate = PyGILState_Ensure();
|
||||
|
||||
PyObject * pyMain = PyImport_AddModule("__main__");
|
||||
|
||||
PyImport_ImportModule("mliap_model_python_couple");
|
||||
if (!pyMain) {
|
||||
PyGILState_Release(gstate);
|
||||
error->all(FLERR,"Could not initialize embedded Python");
|
||||
}
|
||||
|
||||
|
||||
PyObject* coupling_module = PyImport_ImportModule("mliap_model_python_couple");
|
||||
|
||||
if (!coupling_module) {
|
||||
PyErr_Print();
|
||||
PyErr_Clear();
|
||||
PyGILState_Release(gstate);
|
||||
error->all(FLERR,"Loading MLIAPPY coupling module failure.");
|
||||
}
|
||||
|
||||
// Recipe from lammps/src/pair_python.cpp :
|
||||
// add current directory to PYTHONPATH
|
||||
PyObject * py_path = PySys_GetObject((char *)"path");
|
||||
@ -61,7 +58,6 @@ MLIAPModelPython::MLIAPModelPython(LAMMPS* lmp, char* coefffilename) :
|
||||
if (potentials_path != NULL) {
|
||||
PyList_Append(py_path, PY_STRING_FROM_STRING(potentials_path));
|
||||
}
|
||||
|
||||
PyGILState_Release(gstate);
|
||||
|
||||
if (coefffilename) read_coeffs(coefffilename);
|
||||
@ -82,29 +78,51 @@ MLIAPModelPython::~MLIAPModelPython(){
|
||||
|
||||
int MLIAPModelPython::get_nparams()
|
||||
{
|
||||
if (nparams == 0) {
|
||||
if (ndescriptors == 0) error->all(FLERR,"ndescriptors not defined");
|
||||
else nparams = ndescriptors + 1;
|
||||
}
|
||||
return nparams;
|
||||
}
|
||||
|
||||
|
||||
void MLIAPModelPython::read_coeffs(char * fname)
|
||||
{
|
||||
PyGILState_STATE gstate = PyGILState_Ensure();
|
||||
int err = MLIAPPY_load_model(this, fname);
|
||||
if (err) {
|
||||
|
||||
int loaded = MLIAPPY_load_model(this, fname);
|
||||
if (PyErr_Occurred()) {
|
||||
PyErr_Print();
|
||||
PyErr_Clear();
|
||||
PyGILState_Release(gstate);
|
||||
error->all(FLERR,"Loading python model failure.");
|
||||
}
|
||||
PyGILState_Release(gstate);
|
||||
|
||||
if (loaded) {
|
||||
this->connect_param_counts();
|
||||
}
|
||||
else {
|
||||
utils::logmesg(lmp,"Loading python model deferred.\n");
|
||||
}
|
||||
}
|
||||
|
||||
// Finalize loading of the model.
|
||||
void MLIAPModelPython::connect_param_counts()
|
||||
{
|
||||
PyGILState_STATE gstate = PyGILState_Ensure();
|
||||
nelements = MLIAPPY_nelements(this);
|
||||
nparams = MLIAPPY_nparams(this);
|
||||
ndescriptors = MLIAPPY_ndescriptors(this);
|
||||
|
||||
if (PyErr_Occurred()) {
|
||||
PyErr_Print();
|
||||
PyErr_Clear();
|
||||
PyGILState_Release(gstate);
|
||||
error->all(FLERR,"Loading python model failure.");
|
||||
}
|
||||
PyGILState_Release(gstate);
|
||||
model_loaded = 1;
|
||||
utils::logmesg(lmp,"Loading python model complete.\n");
|
||||
|
||||
}
|
||||
|
||||
|
||||
/* ----------------------------------------------------------------------
|
||||
Calculate model gradients w.r.t descriptors
|
||||
for each atom beta_i = dE(B_i)/dB_i
|
||||
@ -112,7 +130,20 @@ void MLIAPModelPython::read_coeffs(char * fname)
|
||||
|
||||
void MLIAPModelPython::compute_gradients(MLIAPData* data)
|
||||
{
|
||||
MLIAPPY_model_callback(this, data);
|
||||
if (not model_loaded) {
|
||||
error->all(FLERR,"Model not loaded.");
|
||||
}
|
||||
|
||||
PyGILState_STATE gstate = PyGILState_Ensure();
|
||||
MLIAPPY_compute_gradients(this, data);
|
||||
if (PyErr_Occurred()) {
|
||||
PyErr_Print();
|
||||
PyErr_Clear();
|
||||
PyGILState_Release(gstate);
|
||||
error->all(FLERR,"Running python model failure.");
|
||||
}
|
||||
PyGILState_Release(gstate);
|
||||
|
||||
}
|
||||
|
||||
/* ----------------------------------------------------------------------
|
||||
|
||||
@ -29,9 +29,14 @@ public:
|
||||
virtual void compute_gradgrads(class MLIAPData*);
|
||||
virtual void compute_force_gradients(class MLIAPData*);
|
||||
virtual double memory_usage();
|
||||
void connect_param_counts(); // If possible convert this to protected/private and
|
||||
// and figure out how to declare cython fn
|
||||
// load_from_python as a friend.
|
||||
int model_loaded;
|
||||
|
||||
protected:
|
||||
virtual void read_coeffs(char *);
|
||||
|
||||
|
||||
private:
|
||||
|
||||
};
|
||||
|
||||
@ -1,5 +1,12 @@
|
||||
# cython: language_level=3
|
||||
# distutils: language = c++
|
||||
# distutils: define_macros="LMP_MLIAPPY"
|
||||
# distutils: extra_compile_args= -stdlib=libc++ -std=c++11
|
||||
# distutils: include_dirs = ../STUBS .. ../MLIAP
|
||||
# distutils: extra_link_args= -stdlib=libc++
|
||||
# Note: only the language_level and language commands are needed, the rest pertain
|
||||
# to building mliap_model_python_couple as a standalone python extension, which
|
||||
# is experimental.
|
||||
|
||||
cimport cython
|
||||
|
||||
@ -7,7 +14,6 @@ import pickle
|
||||
|
||||
# For converting C arrays to numpy arrays
|
||||
import numpy as np
|
||||
cimport numpy as cnp
|
||||
|
||||
# For converting void * to integer for tracking object identity
|
||||
from libc.stdint cimport uintptr_t
|
||||
@ -32,40 +38,72 @@ cdef extern from "mliap_data.h" namespace "LAMMPS_NS":
|
||||
|
||||
cdef extern from "mliap_model_python.h" namespace "LAMMPS_NS":
|
||||
cdef cppclass MLIAPModelPython:
|
||||
ctypedef void (*CBPtr)(void * , MLIAPData);
|
||||
void set_model(CBPtr, void *);
|
||||
|
||||
void connect_param_counts()
|
||||
|
||||
|
||||
class MLIAPPYModelNotLinked(Exception): pass
|
||||
|
||||
|
||||
LOADED_MODELS = {}
|
||||
cdef public int MLIAPPY_load_model(MLIAPModelPython * c_model, char* fname) except 1 with gil:
|
||||
str_fname = fname.decode('utf-8') # Python 3 only; not Python 2 not supported.
|
||||
|
||||
with open(str_fname,'rb') as pfile:
|
||||
model = pickle.load(pfile)
|
||||
|
||||
LOADED_MODELS[int(<uintptr_t> c_model)] = model
|
||||
return 0
|
||||
cdef object c_id(MLIAPModelPython * c_model):
|
||||
"""
|
||||
Use python-style id of object to keep track of identity.
|
||||
Note, this is probably not a perfect general strategy but it should work fine with LAMMPS pair styles.
|
||||
"""
|
||||
return int(<uintptr_t> c_model)
|
||||
|
||||
cdef object retrieve(MLIAPModelPython * c_model):
|
||||
try:
|
||||
model = LOADED_MODELS[c_id(c_model)]
|
||||
except KeyError as ke:
|
||||
raise KeyError("Model has not been loaded.") from ke
|
||||
if model is None:
|
||||
raise MLIAPPYModelNotLinked("Model not linked, connect the model from the python side.")
|
||||
return model
|
||||
|
||||
cdef public int MLIAPPY_load_model(MLIAPModelPython * c_model, char* fname) with gil:
|
||||
str_fname = fname.decode('utf-8') # Python 3 only; not Python 2 not supported.
|
||||
if str_fname == "LATER":
|
||||
model = None
|
||||
returnval = 0
|
||||
else:
|
||||
with open(str_fname,'rb') as pfile:
|
||||
model = pickle.load(pfile)
|
||||
returnval = 1
|
||||
LOADED_MODELS[c_id(c_model)] = model
|
||||
return returnval
|
||||
|
||||
def load_from_python(model):
|
||||
unloaded_models = [k for k, v in LOADED_MODELS.items() if v is None]
|
||||
num_models = len(unloaded_models)
|
||||
cdef MLIAPModelPython * lmp_model
|
||||
|
||||
if num_models == 0:
|
||||
raise ValueError("No model in the waiting area.")
|
||||
elif num_models > 1:
|
||||
raise ValueError("Model is amibguous, more than one model in waiting area.")
|
||||
else:
|
||||
c_id = unloaded_models[0]
|
||||
LOADED_MODELS[c_id]=model
|
||||
lmp_model = <MLIAPModelPython *> <uintptr_t> c_id
|
||||
lmp_model.connect_param_counts()
|
||||
|
||||
|
||||
cdef public void MLIAPPY_unload_model(MLIAPModelPython * c_model) with gil:
|
||||
del LOADED_MODELS[int(<uintptr_t> c_model)]
|
||||
del LOADED_MODELS[c_id(c_model)]
|
||||
|
||||
cdef public int MLIAPPY_nparams(MLIAPModelPython * c_model) with gil:
|
||||
model = LOADED_MODELS[int(<uintptr_t> c_model)]
|
||||
n_params = int(model.n_params)
|
||||
return <int> n_params
|
||||
return int(retrieve(c_model).n_params)
|
||||
|
||||
cdef public int MLIAPPY_nelements(MLIAPModelPython * c_model) with gil:
|
||||
model = LOADED_MODELS[int(<uintptr_t> c_model)]
|
||||
n_elements = int(model.n_elements)
|
||||
return <int> n_elements
|
||||
return int(retrieve(c_model).n_elements)
|
||||
|
||||
cdef public int MLIAPPY_ndescriptors(MLIAPModelPython * c_model) with gil:
|
||||
model = LOADED_MODELS[int(<uintptr_t> c_model)]
|
||||
n_descriptors = int(model.n_descriptors)
|
||||
return <int> n_descriptors
|
||||
return int(retrieve(c_model).n_descriptors)
|
||||
|
||||
cdef public MLIAPPY_model_callback(MLIAPModelPython * c_model, MLIAPData * data) with gil:
|
||||
model = LOADED_MODELS[int(<uintptr_t> c_model)]
|
||||
cdef public void MLIAPPY_compute_gradients(MLIAPModelPython * c_model, MLIAPData * data) with gil:
|
||||
model = retrieve(c_model)
|
||||
|
||||
n_d = data.ndescriptors
|
||||
n_a = data.natoms
|
||||
@ -73,11 +111,11 @@ cdef public MLIAPPY_model_callback(MLIAPModelPython * c_model, MLIAPData * data)
|
||||
# Make numpy arrays from pointers
|
||||
beta_np = np.asarray(<double[:n_a,:n_d] > &data.betas[0][0])
|
||||
desc_np = np.asarray(<double[:n_a,:n_d]> &data.descriptors[0][0])
|
||||
type_np = np.asarray(<int[:n_a]> &data.ielems[0])
|
||||
elem_np = np.asarray(<int[:n_a]> &data.ielems[0])
|
||||
en_np = np.asarray(<double[:n_a]> &data.eatoms[0])
|
||||
|
||||
# Invoke python model on numpy arrays.
|
||||
model(type_np,desc_np,beta_np,en_np)
|
||||
model(elem_np,desc_np,beta_np,en_np)
|
||||
|
||||
# Get the total energy from the atom energy.
|
||||
energy = np.sum(en_np)
|
||||
|
||||
@ -15,14 +15,14 @@ class TorchWrapper(torch.nn.Module):
|
||||
self.n_descriptors = n_descriptors
|
||||
self.n_elements = n_elements
|
||||
|
||||
def __call__(self, types, bispectrum, beta, energy):
|
||||
def __call__(self, elems, bispectrum, beta, energy):
|
||||
|
||||
bispectrum = torch.from_numpy(bispectrum).to(self.dtype).requires_grad_(True)
|
||||
types = torch.from_numpy(types).to(torch.long) - 1
|
||||
elems = torch.from_numpy(elems).to(torch.long) - 1
|
||||
|
||||
with torch.autograd.enable_grad():
|
||||
|
||||
energy_nn = self.model(bispectrum, types)
|
||||
energy_nn = self.model(bispectrum, elems)
|
||||
if energy_nn.ndim > 1:
|
||||
energy_nn = energy_nn.flatten()
|
||||
|
||||
@ -37,10 +37,10 @@ class TorchWrapper32(TorchWrapper):
|
||||
class TorchWrapper64(TorchWrapper):
|
||||
dtype = torch.float64
|
||||
|
||||
class IgnoreTypes(torch.nn.Module):
|
||||
class IgnoreElems(torch.nn.Module):
|
||||
def __init__(self,subnet):
|
||||
super().__init__()
|
||||
self.subnet = subnet
|
||||
|
||||
def forward(self,bispectrum,types):
|
||||
def forward(self,bispectrum,elems):
|
||||
return self.subnet(bispectrum)
|
||||
|
||||
@ -28,6 +28,9 @@
|
||||
|
||||
#ifdef LMP_MLIAPPY
|
||||
#include "mliap_model_python.h"
|
||||
// The above should somehow really be included in the next file.
|
||||
// We could get around this with cython --capi-reexport-cincludes
|
||||
// However, that exposes -too many- headers.
|
||||
#include "mliap_model_python_couple.h"
|
||||
#endif
|
||||
|
||||
@ -52,7 +55,6 @@ PythonImpl::PythonImpl(LAMMPS *lmp) : Pointers(lmp)
|
||||
|
||||
nfunc = 0;
|
||||
pfuncs = nullptr;
|
||||
|
||||
// one-time initialization of Python interpreter
|
||||
// pyMain stores pointer to main module
|
||||
external_interpreter = Py_IsInitialized();
|
||||
@ -61,7 +63,7 @@ PythonImpl::PythonImpl(LAMMPS *lmp) : Pointers(lmp)
|
||||
// Inform python intialization scheme of the mliappy module.
|
||||
// This -must- happen before python is initialized.
|
||||
int err = PyImport_AppendInittab("mliap_model_python_couple", PyInit_mliap_model_python_couple);
|
||||
// todo: catch if error and report problem.
|
||||
if (err) error->all(FLERR,"Could not register MLIAPPY embedded python module.");
|
||||
#endif
|
||||
|
||||
Py_Initialize();
|
||||
|
||||
Reference in New Issue
Block a user