diff --git a/.gitignore b/.gitignore index 0f1b01775d..92fae19766 100644 --- a/.gitignore +++ b/.gitignore @@ -46,3 +46,4 @@ Thumbs.db /Makefile /cmake_install.cmake /lmp + diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 42e6d12ffb..fab10bfdef 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -105,7 +105,7 @@ install(TARGETS lmp EXPORT LAMMPS_Targets DESTINATION ${CMAKE_INSTALL_BINDIR}) option(CMAKE_VERBOSE_MAKEFILE "Generate verbose Makefiles" OFF) set(STANDARD_PACKAGES ASPHERE BODY CLASS2 COLLOID COMPRESS DIPOLE - GRANULAR KSPACE LATTE MANYBODY MC MESSAGE MISC MLIAP MOLECULE PERI POEMS + GRANULAR KSPACE LATTE MANYBODY MC MESSAGE MISC MLIAP MLIAPPY MOLECULE PERI POEMS QEQ REPLICA RIGID SHOCK SPIN SNAP SRD KIM PYTHON MSCG MPIIO VORONOI USER-ATC USER-AWPMD USER-BOCS USER-CGDNA USER-MESODPD USER-CGSDK USER-COLVARS USER-DIFFRACTION USER-DPD USER-DRUDE USER-EFF USER-FEP USER-H5MD USER-LB @@ -194,6 +194,8 @@ endif() # "hard" dependencies between packages resulting # in an error instead of skipping over files pkg_depends(MLIAP SNAP) +pkg_depends(MLIAPPY MLIAP) +pkg_depends(MLIAPPY PYTHON) pkg_depends(MPIIO MPI) pkg_depends(USER-ATC MANYBODY) pkg_depends(USER-LB MPI) @@ -372,7 +374,7 @@ else() set(CUDA_REQUEST_PIC) endif() -foreach(PKG_WITH_INCL KSPACE PYTHON VORONOI USER-COLVARS USER-MOLFILE USER-NETCDF USER-PLUMED USER-QMMM +foreach(PKG_WITH_INCL KSPACE PYTHON MLIAPPY VORONOI USER-COLVARS USER-MOLFILE USER-NETCDF USER-PLUMED USER-QMMM USER-QUIP USER-SCAFACOS USER-SMD USER-VTK KIM LATTE MESSAGE MSCG COMPRESS) if(PKG_${PKG_WITH_INCL}) include(Packages/${PKG_WITH_INCL}) diff --git a/cmake/Modules/Packages/MLIAPPY.cmake b/cmake/Modules/Packages/MLIAPPY.cmake new file mode 100644 index 0000000000..73ce83361c --- /dev/null +++ b/cmake/Modules/Packages/MLIAPPY.cmake @@ -0,0 +1,3 @@ +execute_process(COMMAND cythonize mliap_model_python_couple.pyx WORKING_DIRECTORY ../src/MLIAPPY) + +target_compile_definitions(lammps PRIVATE -DLMP_MLIAPPY) diff --git a/cmake/Modules/Packages/PYTHON.cmake b/cmake/Modules/Packages/PYTHON.cmake index 7be25a6b05..e3158dc509 100644 --- a/cmake/Modules/Packages/PYTHON.cmake +++ b/cmake/Modules/Packages/PYTHON.cmake @@ -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) diff --git a/doc/src/Packages_details.rst b/doc/src/Packages_details.rst index 2b2cb90aff..f8e18dafbc 100644 --- a/doc/src/Packages_details.rst +++ b/doc/src/Packages_details.rst @@ -45,6 +45,7 @@ page gives those details. * :ref:`MESSAGE ` * :ref:`MISC ` * :ref:`MLIAP ` + * :ref:`MLIAPPY ` * :ref:`MOLECULE ` * :ref:`MPIIO ` * :ref:`MSCG ` @@ -678,6 +679,36 @@ To use this package, also the :ref:`SNAP package ` needs to be install ---------- +.. _PKG-MLIAPPY: + +MLIAPPY package +------------- + +**Contents:** + +Extension to the MLIAP package for coupling with python models. + +**Install:** + +To use this package, also the :ref:`MLIAP package ` needs to be installed. +To use this package, also the :ref:`PYTHON package ` needs to be installed. +The version of python must be >3.5. + +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. + +**Author:** Nicholas Lubbers (LANL). + +**Supporting info:** + +* src/MLIAPPY: filenames -> commands +* src/MLIAPPY/README +* :doc:`pair_style mliap ` +* examples/mliappy (see README) +---------- + .. _PKG-MOLECULE: MOLECULE package diff --git a/doc/src/Packages_standard.rst b/doc/src/Packages_standard.rst index ab9be69ab3..800fd3c501 100644 --- a/doc/src/Packages_standard.rst +++ b/doc/src/Packages_standard.rst @@ -61,6 +61,8 @@ package: +----------------------------------+--------------------------------------+----------------------------------------------------+------------------------------------------------------+---------+ | :ref:`MLIAP ` | multiple machine learning potentials | :doc:`pair_style mliap ` | mliap | no | +----------------------------------+--------------------------------------+----------------------------------------------------+------------------------------------------------------+---------+ +| :ref:`MLIAPPY ` | enable python ML models for mliap | :doc:`pair_style mliap ` | mliap | no | ++----------------------------------+--------------------------------------+----------------------------------------------------+------------------------------------------------------+---------+ | :ref:`MOLECULE ` | molecular system force fields | :doc:`Howto bioFF ` | peptide | no | +----------------------------------+--------------------------------------+----------------------------------------------------+------------------------------------------------------+---------+ | :ref:`MPIIO ` | MPI parallel I/O dump and restart | :doc:`dump ` | n/a | no | diff --git a/doc/src/compute_mliap.rst b/doc/src/compute_mliap.rst index fcc336e682..e4e6603cdc 100644 --- a/doc/src/compute_mliap.rst +++ b/doc/src/compute_mliap.rst @@ -18,7 +18,7 @@ Syntax .. parsed-literal:: *model* values = style - style = *linear* or *quadratic* + style = *linear* or *quadratic* or *mliappy* *descriptor* values = style filename style = *sna* filename = name of file containing descriptor definitions @@ -57,7 +57,8 @@ The compute *mliap* command must be followed by two keywords *model* and *descriptor* in either order. The *model* keyword is followed by a model style, currently limited to -either *linear* or *quadratic*. +either *linear* or *quadratic*. The *mliappy* model is only available +if lammps is built with MLIAPPY package. The *descriptor* keyword is followed by a descriptor style, and additional arguments. Currently the only descriptor style is *sna*, indicating the bispectrum component @@ -167,6 +168,8 @@ LAMMPS was built with that package. In addition, building LAMMPS with the MLIAP requires building LAMMPS with the SNAP package. See the :doc:`Build package ` doc page for more info. +Python models such as neural networks can be used if the MLIAPPY package is built. + Related commands """""""""""""""" diff --git a/doc/src/pair_mliap.rst b/doc/src/pair_mliap.rst index 09295cd8be..aeba2d9fd2 100644 --- a/doc/src/pair_mliap.rst +++ b/doc/src/pair_mliap.rst @@ -16,7 +16,7 @@ Syntax .. parsed-literal:: *model* values = style filename - style = *linear* or *quadratic* + style = *linear* or *quadratic* or *mliappy* filename = name of file containing model definitions *descriptor* values = style filename style = *sna* @@ -43,9 +43,10 @@ it is possible to use many different models with a given descriptor, or many different descriptors with a given model. Currently, the pair_style supports just two models, *linear* and *quadratic*, and one descriptor, *sna*, the SNAP descriptor used by :doc:`pair_style snap `, including the linear, quadratic, -and chem variants. Work is currently underway to extend -the interface to handle neural network energy models, -and it is also straightforward to add new descriptor styles. +and chem variants. With the MLIAPPY package installed, the *mliappy* model +is available which can be used to couple python models such as neural network +energy models. +It is also straightforward to add new descriptor styles. In order to train a model, it is useful to know the gradient or derivative of energy, force, and stress w.r.t. model parameters. This information can be accessed using the related :doc:`compute mliap ` command. @@ -143,6 +144,9 @@ was built with that package. In addition, building LAMMPS with the MLIAP package requires building LAMMPS with the SNAP package. See the :doc:`Build package ` doc page for more info. +Python models such as neural networks can be used if the MLIAPPY package is built. + + Related commands """""""""""""""" diff --git a/examples/mliappy/README b/examples/mliappy/README new file mode 100644 index 0000000000..d61972afcf --- /dev/null +++ b/examples/mliappy/README @@ -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` diff --git a/examples/mliappy/Ta06A.mliap.descriptor b/examples/mliappy/Ta06A.mliap.descriptor new file mode 100644 index 0000000000..481ebf7e44 --- /dev/null +++ b/examples/mliappy/Ta06A.mliap.descriptor @@ -0,0 +1,21 @@ +# DATE: 2014-09-05 UNITS: metal CONTRIBUTOR: Aidan Thompson athomps@sandia.gov CITATION: Thompson, Swiler, Trott, Foiles and Tucker, arxiv.org, 1409.3880 (2014) + +# LAMMPS SNAP parameters for Ta_Cand06A + +# required +rcutfac 4.67637 +twojmax 6 + +# elements + +nelems 1 +elems Ta +radelems 0.5 +welems 1 + +# optional + +rfac0 0.99363 +rmin0 0 +bzeroflag 0 + diff --git a/examples/mliappy/Ta06A.mliap.pytorch b/examples/mliappy/Ta06A.mliap.pytorch new file mode 100644 index 0000000000..c6c6bc653d --- /dev/null +++ b/examples/mliappy/Ta06A.mliap.pytorch @@ -0,0 +1,18 @@ +# 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 Ta06A.mliap.pytorch.model.pkl & +descriptor sna Ta06A.mliap.descriptor +pair_coeff 1 1 zbl ${zblz} ${zblz} +pair_coeff * * mliap Ta + diff --git a/examples/mliappy/convert_mliap_Ta06A.py b/examples/mliappy/convert_mliap_Ta06A.py new file mode 100644 index 0000000000..371139481d --- /dev/null +++ b/examples/mliappy/convert_mliap_Ta06A.py @@ -0,0 +1,33 @@ +import sys +import numpy as np +import torch +import pickle +import os +import shutil + +shutil.copyfile('../../src/MLIAPPY/mliappy_pytorch.py','./mliappy_pytorch.py') + +import mliappy_pytorch + +# Read coefficients +coeffs = np.genfromtxt("../mliap/Ta06A.mliap.model",skip_header=6) + +# Write coefficiets to a pytorch linear model +bias = coeffs[0] +weights = coeffs[1:] +lin = torch.nn.Linear(weights.shape[0],1) +lin.to(torch.float64) +with torch.autograd.no_grad(): + lin.weight.set_(torch.from_numpy(weights).unsqueeze(0)) + lin.bias.set_(torch.as_tensor(bias,dtype=torch.float64).unsqueeze(0)) + +# Wrap the pytorch model for usage with MLIAPPY +model = mliappy_pytorch.IgnoreElems(lin) +n_descriptors = lin.weight.shape[1] +n_params = mliappy_pytorch.calc_n_params(model) +n_types = 1 +linked_model = mliappy_pytorch.TorchWrapper64(model,n_descriptors=n_descriptors,n_elements=n_types) + +# Save the result. +with open("Ta06A.mliap.pytorch.model.pkl",'wb') as pfile: + pickle.dump(linked_model,pfile) diff --git a/examples/mliappy/in.mliap.pytorch.Ta06A b/examples/mliappy/in.mliap.pytorch.Ta06A new file mode 100644 index 0000000000..1d93d6a424 --- /dev/null +++ b/examples/mliappy/in.mliap.pytorch.Ta06A @@ -0,0 +1,53 @@ +# 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 + +include Ta06A.mliap.pytorch + +# 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} + diff --git a/examples/mliappy/load_external.py b/examples/mliappy/load_external.py new file mode 100644 index 0000000000..606ee8f940 --- /dev/null +++ b/examples/mliappy/load_external.py @@ -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) + + diff --git a/examples/mliappy/test_pylibs.py b/examples/mliappy/test_pylibs.py new file mode 100644 index 0000000000..c5dcab75cf --- /dev/null +++ b/examples/mliappy/test_pylibs.py @@ -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") diff --git a/python/lammps.py b/python/lammps.py index b74e111dc6..85fbd1217b 100644 --- a/python/lammps.py +++ b/python/lammps.py @@ -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("") + +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) + diff --git a/src/MLIAP/compute_mliap.cpp b/src/MLIAP/compute_mliap.cpp index 34d086462f..64753fe64f 100644 --- a/src/MLIAP/compute_mliap.cpp +++ b/src/MLIAP/compute_mliap.cpp @@ -17,6 +17,10 @@ #include "mliap_model_linear.h" #include "mliap_model_quadratic.h" #include "mliap_descriptor_snap.h" +#ifdef LMP_MLIAPPY +#include "mliap_model_python.h" +#endif + #include "compute_mliap.h" #include "atom.h" #include "update.h" @@ -65,7 +69,14 @@ ComputeMLIAP::ComputeMLIAP(LAMMPS *lmp, int narg, char **arg) : } else if (strcmp(arg[iarg+1],"quadratic") == 0) { model = new MLIAPModelQuadratic(lmp); iarg += 2; - } else error->all(FLERR,"Illegal compute mliap command"); + } + #ifdef LMP_MLIAPPY + else if (strcmp(arg[iarg+1],"mliappy") == 0) { + model = new MLIAPModelPython(lmp); + iarg += 2; + } + #endif + else error->all(FLERR,"Illegal compute mliap command"); modelflag = 1; } else if (strcmp(arg[iarg],"descriptor") == 0) { if (iarg+2 > narg) error->all(FLERR,"Illegal compute mliap command"); diff --git a/src/MLIAP/mliap_data.cpp b/src/MLIAP/mliap_data.cpp index e6502d001a..3ada98e8c0 100644 --- a/src/MLIAP/mliap_data.cpp +++ b/src/MLIAP/mliap_data.cpp @@ -25,7 +25,8 @@ MLIAPData::MLIAPData(LAMMPS *lmp, int gradgradflag_in, int *map_in, class MLIAPModel* model_in, class MLIAPDescriptor* descriptor_in, class PairMLIAP* pairmliap_in) : - Pointers(lmp), gradforce(nullptr), betas(nullptr), descriptors(nullptr), gamma(nullptr), + Pointers(lmp), gradforce(nullptr), betas(nullptr), + descriptors(nullptr), eatoms(nullptr), gamma(nullptr), gamma_row_index(nullptr), gamma_col_index(nullptr), egradient(nullptr), numneighs(nullptr), iatoms(nullptr), ielems(nullptr), jatoms(nullptr), jelems(nullptr), rij(nullptr), graddesc(nullptr), model(nullptr), descriptor(nullptr), list(nullptr) @@ -64,6 +65,7 @@ MLIAPData::~MLIAPData() { memory->destroy(betas); memory->destroy(descriptors); + memory->destroy(eatoms); memory->destroy(gamma_row_index); memory->destroy(gamma_col_index); memory->destroy(gamma); @@ -133,6 +135,7 @@ void MLIAPData::generate_neighdata(NeighList* list_in, int eflag_in, int vflag_i if (natoms_max < natoms) { memory->grow(betas,natoms,ndescriptors,"MLIAPData:betas"); memory->grow(descriptors,natoms,ndescriptors,"MLIAPData:descriptors"); + memory->grow(eatoms,natoms,"MLIAPData:eatoms"); natoms_max = natoms; } @@ -267,6 +270,7 @@ double MLIAPData::memory_usage() bytes += natoms*ndescriptors*sizeof(int); // betas bytes += natoms*ndescriptors*sizeof(int); // descriptors + bytes += natoms*sizeof(double); // eatoms bytes += natomneigh_max*sizeof(int); // iatoms bytes += natomneigh_max*sizeof(int); // ielems diff --git a/src/MLIAP/mliap_data.h b/src/MLIAP/mliap_data.h index ffac9ccd4c..7fb60bd723 100644 --- a/src/MLIAP/mliap_data.h +++ b/src/MLIAP/mliap_data.h @@ -34,8 +34,10 @@ class MLIAPData : protected Pointers { int yoffset, zoffset; int ndims_force, ndims_virial; double **gradforce; - double** betas; // betas for all atoms in list + double** betas; // betas for all atoms in list double** descriptors; // descriptors for all atoms in list + double* eatoms; // energies for all atoms in list + double energy; // energy int ndescriptors; // number of descriptors int nparams; // number of model parameters per element int nelements; // number of elements diff --git a/src/MLIAP/mliap_model.cpp b/src/MLIAP/mliap_model.cpp index 04b2ac663a..a4f5cdc86c 100644 --- a/src/MLIAP/mliap_model.cpp +++ b/src/MLIAP/mliap_model.cpp @@ -28,13 +28,9 @@ using namespace LAMMPS_NS; MLIAPModel::MLIAPModel(LAMMPS* lmp, char* coefffilename) : Pointers(lmp) { - coeffelem = nullptr; - if (coefffilename) read_coeffs(coefffilename); - else { - nparams = 0; - nelements = 0; - ndescriptors = 0; - } + nparams = 0; + nelements = 0; + ndescriptors = 0; nonlinearflag = 0; } @@ -42,9 +38,9 @@ MLIAPModel::MLIAPModel(LAMMPS* lmp, char* coefffilename) : Pointers(lmp) MLIAPModel::~MLIAPModel() { - memory->destroy(coeffelem); } + /* ---------------------------------------------------------------------- placeholder ------------------------------------------------------------------------- */ @@ -73,7 +69,22 @@ void MLIAPModel::set_ndescriptors(int ndescriptors_in) /* ---------------------------------------------------------------------- */ -void MLIAPModel::read_coeffs(char *coefffilename) + +MLIAPModelSimple::MLIAPModelSimple(LAMMPS* lmp, char* coefffilename) : MLIAPModel(lmp, coefffilename) +{ + coeffelem = nullptr; + if (coefffilename) read_coeffs(coefffilename); +} + +/* ---------------------------------------------------------------------- */ + +MLIAPModelSimple::~MLIAPModelSimple() +{ + memory->destroy(coeffelem); +} + + +void MLIAPModelSimple::read_coeffs(char *coefffilename) { // open coefficient file on proc 0 @@ -165,7 +176,7 @@ void MLIAPModel::read_coeffs(char *coefffilename) memory usage ------------------------------------------------------------------------- */ -double MLIAPModel::memory_usage() +double MLIAPModelSimple::memory_usage() { double bytes = 0; diff --git a/src/MLIAP/mliap_model.h b/src/MLIAP/mliap_model.h index b2f7312897..96cfff0a3d 100644 --- a/src/MLIAP/mliap_model.h +++ b/src/MLIAP/mliap_model.h @@ -30,15 +30,26 @@ public: virtual void compute_gradgrads(class MLIAPData*)=0; virtual void compute_force_gradients(class MLIAPData*)=0; virtual void init(); - virtual double memory_usage(); + virtual double memory_usage()=0; int nelements; // # of unique elements - int nonlinearflag; // 1 if gradient() requires escriptors + int nonlinearflag; // 1 if gradient() requires descriptors int ndescriptors; // number of descriptors int nparams; // number of parameters per element protected: - void read_coeffs(char *); + virtual void read_coeffs(char *)=0; +}; + +class MLIAPModelSimple : public MLIAPModel { +public: + MLIAPModelSimple(LAMMPS*, char*); + ~MLIAPModelSimple(); + virtual double memory_usage(); + +protected: double **coeffelem; // element coefficients + virtual void read_coeffs(char *); + }; } diff --git a/src/MLIAP/mliap_model_linear.cpp b/src/MLIAP/mliap_model_linear.cpp index b3224284e5..428d163166 100644 --- a/src/MLIAP/mliap_model_linear.cpp +++ b/src/MLIAP/mliap_model_linear.cpp @@ -21,7 +21,7 @@ using namespace LAMMPS_NS; /* ---------------------------------------------------------------------- */ MLIAPModelLinear::MLIAPModelLinear(LAMMPS* lmp, char* coefffilename) : - MLIAPModel(lmp, coefffilename) + MLIAPModelSimple(lmp, coefffilename) { if (nparams > 0) ndescriptors = nparams - 1; } @@ -51,8 +51,9 @@ int MLIAPModelLinear::get_nparams() void MLIAPModelLinear::compute_gradients(MLIAPData* data) { + data->energy = 0.0; + for (int ii = 0; ii < data->natoms; ii++) { - const int i = data->iatoms[ii]; const int ielem = data->ielems[ii]; double* coeffi = coeffelem[ielem]; @@ -74,7 +75,8 @@ void MLIAPModelLinear::compute_gradients(MLIAPData* data) for (int icoeff = 0; icoeff < data->ndescriptors; icoeff++) etmp += coeffi[icoeff+1]*data->descriptors[ii][icoeff]; - data->pairmliap->e_tally(i,etmp); + data->energy += etmp; + data->eatoms[ii] = etmp; } } } diff --git a/src/MLIAP/mliap_model_linear.h b/src/MLIAP/mliap_model_linear.h index 326407faa8..89c69d3537 100644 --- a/src/MLIAP/mliap_model_linear.h +++ b/src/MLIAP/mliap_model_linear.h @@ -18,7 +18,7 @@ namespace LAMMPS_NS { -class MLIAPModelLinear : public MLIAPModel { +class MLIAPModelLinear : public MLIAPModelSimple { public: MLIAPModelLinear(LAMMPS*, char* = nullptr); ~MLIAPModelLinear(); diff --git a/src/MLIAP/mliap_model_quadratic.cpp b/src/MLIAP/mliap_model_quadratic.cpp index 481417078a..222f6ad0ae 100644 --- a/src/MLIAP/mliap_model_quadratic.cpp +++ b/src/MLIAP/mliap_model_quadratic.cpp @@ -22,8 +22,9 @@ using namespace LAMMPS_NS; /* ---------------------------------------------------------------------- */ MLIAPModelQuadratic::MLIAPModelQuadratic(LAMMPS* lmp, char* coefffilename) : - MLIAPModel(lmp, coefffilename) + MLIAPModelSimple(lmp, coefffilename) { + if (coefffilename) read_coeffs(coefffilename); if (nparams > 0) ndescriptors = sqrt(2*nparams)-1; nonlinearflag = 1; } @@ -52,8 +53,9 @@ int MLIAPModelQuadratic::get_nparams() void MLIAPModelQuadratic::compute_gradients(MLIAPData* data) { + data->energy = 0.0; + for (int ii = 0; ii < data->natoms; ii++) { - const int i = data->iatoms[ii]; const int ielem = data->ielems[ii]; double* coeffi = coeffelem[ielem]; @@ -99,7 +101,8 @@ void MLIAPModelQuadratic::compute_gradients(MLIAPData* data) etmp += coeffi[k++]*bveci*bvecj; } } - data->pairmliap->e_tally(i,etmp); + data->energy += etmp; + data->eatoms[ii] = etmp; } } } diff --git a/src/MLIAP/mliap_model_quadratic.h b/src/MLIAP/mliap_model_quadratic.h index b41df18a07..28f1732f9b 100644 --- a/src/MLIAP/mliap_model_quadratic.h +++ b/src/MLIAP/mliap_model_quadratic.h @@ -18,7 +18,7 @@ namespace LAMMPS_NS { -class MLIAPModelQuadratic : public MLIAPModel { +class MLIAPModelQuadratic : public MLIAPModelSimple { public: MLIAPModelQuadratic(LAMMPS*, char* = nullptr); ~MLIAPModelQuadratic(); diff --git a/src/MLIAP/pair_mliap.cpp b/src/MLIAP/pair_mliap.cpp index a422271f19..7a2d5b14b5 100644 --- a/src/MLIAP/pair_mliap.cpp +++ b/src/MLIAP/pair_mliap.cpp @@ -17,6 +17,9 @@ #include "mliap_model_linear.h" #include "mliap_model_quadratic.h" #include "mliap_descriptor_snap.h" +#ifdef LMP_MLIAPPY +#include "mliap_model_python.h" +#endif #include "atom.h" #include "error.h" @@ -27,6 +30,7 @@ #include #include +#include "error.h" using namespace LAMMPS_NS; @@ -65,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); @@ -78,6 +93,8 @@ void PairMLIAP::compute(int eflag, int vflag) model->compute_gradients(data); + e_tally(data); + // calculate force contributions beta_i*dB_i/dR_j descriptor->compute_forces(data); @@ -107,6 +124,7 @@ void PairMLIAP::allocate() void PairMLIAP::settings(int narg, char ** arg) { + if (narg < 4) error->all(FLERR,"Illegal pair_style command"); @@ -130,6 +148,12 @@ 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) { + if (iarg+3 > narg) error->all(FLERR,"Illegal pair_style mliap command"); + model = new MLIAPModelPython(lmp,arg[iarg+2]); + iarg += 3; + #endif } else error->all(FLERR,"Illegal pair_style mliap command"); modelflag = 1; } else if (strcmp(arg[iarg],"descriptor") == 0) { @@ -211,22 +235,21 @@ 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"); - if (data->nelements != model->nelements) - error->all(FLERR,"Incompatible model and descriptor definitions"); } /* ---------------------------------------------------------------------- - add energy of atom i to global and per-atom energy + add energies to eng_vdwl and per-atom energy ------------------------------------------------------------------------- */ -void PairMLIAP::e_tally(int i, double ei) +void PairMLIAP::e_tally(MLIAPData* data) { - if (eflag_global) eng_vdwl += ei; - if (eflag_atom) eatom[i] += ei; + if (eflag_global) eng_vdwl += data->energy; + if (eflag_atom) + for (int ii = 0; ii < data->natoms; ii++) { + const int i = data->iatoms[ii]; + eatom[i] += data->eatoms[ii]; + } } /* ---------------------------------------------------------------------- diff --git a/src/MLIAP/pair_mliap.h b/src/MLIAP/pair_mliap.h index 61acbbb278..c31634e923 100644 --- a/src/MLIAP/pair_mliap.h +++ b/src/MLIAP/pair_mliap.h @@ -31,7 +31,7 @@ public: virtual void compute(int, int); void settings(int, char **); virtual void coeff(int, char **); - void e_tally(int, double); + void e_tally(class MLIAPData*); void v_tally(int, int, double*, double*); virtual void init_style(); virtual double init_one(int, int); diff --git a/src/MLIAPPY/Install.sh b/src/MLIAPPY/Install.sh new file mode 100755 index 0000000000..8fccbc2a9c --- /dev/null +++ b/src/MLIAPPY/Install.sh @@ -0,0 +1,53 @@ +# Install/unInstall package files in LAMMPS +# mode = 0/1/2 for uninstall/install/update + +mode=$1 + +# arg1 = file, arg2 = file it depends on + +# enforce using portable C locale +LC_ALL=C +export LC_ALL + +action () { + if (test $mode = 0) then + rm -f ../$1 + elif (! cmp -s $1 ../$1) then + if (test -z "$2" || test -e ../$2) then + cp $1 .. + if (test $mode = 2) then + echo " updating src/$1" + fi + fi + elif (test -n "$2") then + if (test ! -e ../$2) then + rm -f ../$1 + fi + fi +} + +# force rebuild of files using python header + +touch ../lmppython.h + +# all package files with no dependencies + +for file in *.cpp *.h; do + test -f ${file} && action $file +done + +# edit 2 Makefile.package files to include/exclude package info + +if (test $1 = 1) then + + if (test -e ../Makefile.package) then + sed -i -e 's|^PKG_INC =[ \t]*|&-DLMP_MLIAPPY |' ../Makefile.package + fi + +elif (test $1 = 0) then + + if (test -e ../Makefile.package) then + sed -i -e 's/[^ \t]*MLIAPPY[^ \t]* //g' ../Makefile.package + fi + +fi diff --git a/src/MLIAPPY/Makefile.lammps b/src/MLIAPPY/Makefile.lammps new file mode 100644 index 0000000000..f9bf1abc13 --- /dev/null +++ b/src/MLIAPPY/Makefile.lammps @@ -0,0 +1 @@ +#TODO diff --git a/src/MLIAPPY/mliap_model_python.cpp b/src/MLIAPPY/mliap_model_python.cpp new file mode 100644 index 0000000000..1ac61f08b0 --- /dev/null +++ b/src/MLIAPPY/mliap_model_python.cpp @@ -0,0 +1,194 @@ +/* ---------------------------------------------------------------------- + LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator + http://lammps.sandia.gov, Sandia National Laboratories + Steve Plimpton, sjplimp@sandia.gov + + Copyright (2003) Sandia Corporation. Under the terms of Contract + DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains + certain rights in this software. This software is distributed under + the GNU General Public License. + + See the README file in the top-level LAMMPS directory. +------------------------------------------------------------------------- */ + +#include +#include "mliap_model_python.h" +#include "mliap_model_python_couple.h" +#include "pair_mliap.h" +#include "mliap_data.h" +#include "error.h" +#include "utils.h" +#include "lmppython.h" +#include "python_compat.h" + + +using namespace LAMMPS_NS; + +/* ---------------------------------------------------------------------- */ + +MLIAPModelPython::MLIAPModelPython(LAMMPS* lmp, char* coefffilename) : + MLIAPModel(lmp, coefffilename) +{ + model_loaded = 0; + python->init(); + PyGILState_STATE gstate = PyGILState_Ensure(); + + PyObject * pyMain = PyImport_AddModule("__main__"); + + 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"); + PyList_Append(py_path, PY_STRING_FROM_STRING(".")); + + // if LAMMPS_POTENTIALS environment variable is set, add it to PYTHONPATH as well + const char * potentials_path = getenv("LAMMPS_POTENTIALS"); + if (potentials_path != NULL) { + PyList_Append(py_path, PY_STRING_FROM_STRING(potentials_path)); + } + PyGILState_Release(gstate); + + if (coefffilename) read_coeffs(coefffilename); + + nonlinearflag=1; +} + +/* ---------------------------------------------------------------------- */ + +MLIAPModelPython::~MLIAPModelPython(){ + MLIAPPY_unload_model(this); +} + +/* ---------------------------------------------------------------------- + get number of parameters + ---------------------------------------------------------------------- */ + + +int MLIAPModelPython::get_nparams() +{ + return nparams; +} + +void MLIAPModelPython::read_coeffs(char * fname) +{ + PyGILState_STATE gstate = PyGILState_Ensure(); + + 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 + ---------------------------------------------------------------------- */ + +void MLIAPModelPython::compute_gradients(MLIAPData* 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); + +} + +/* ---------------------------------------------------------------------- + Calculate model double gradients w.r.t descriptors and parameters + for each atom energy gamma_lk = d2E(B)/dB_k/dsigma_l, + where sigma_l is a parameter, B_k a descriptor, + and atom subscript i is omitted + + gamma is in CSR format: + nnz = number of non-zero values + gamma_row_index[inz] = l indices, 0 <= l < nparams + gamma_col_indexiinz] = k indices, 0 <= k < ndescriptors + gamma[i][inz] = non-zero values, 0 <= inz < nnz + + egradient is derivative of energy w.r.t. parameters + ---------------------------------------------------------------------- */ + +void MLIAPModelPython::compute_gradgrads(class MLIAPData* data) +{ + error->all(FLERR,"compute_gradgrads not implemented"); +} + +/* ---------------------------------------------------------------------- + calculate gradients of forces w.r.t. parameters + egradient is derivative of energy w.r.t. parameters + ---------------------------------------------------------------------- */ + +void MLIAPModelPython::compute_force_gradients(class MLIAPData* data) +{ + error->all(FLERR,"compute_force_gradients not implemented"); +} + +/* ---------------------------------------------------------------------- + count the number of non-zero entries in gamma matrix + ---------------------------------------------------------------------- */ + +int MLIAPModelPython::get_gamma_nnz(class MLIAPData* data) +{ + // todo: get_gamma_nnz + return 0; +} + + +double MLIAPModelPython::memory_usage() +{ + // todo: get approximate memory usage in coupling code. + return 0; +} diff --git a/src/MLIAPPY/mliap_model_python.h b/src/MLIAPPY/mliap_model_python.h new file mode 100644 index 0000000000..0ef87f86fd --- /dev/null +++ b/src/MLIAPPY/mliap_model_python.h @@ -0,0 +1,47 @@ +/* -*- c++ -*- ---------------------------------------------------------- + LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator + http://lammps.sandia.gov, Sandia National Laboratories + Steve Plimpton, sjplimp@sandia.gov + + Copyright (2003) Sandia Corporation. Under the terms of Contract + DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains + certain rights in this software. This software is distributed under + the GNU General Public License. + + See the README file in the top-level LAMMPS directory. +------------------------------------------------------------------------- */ + +#ifndef LMP_MLIAP_MODEL_PYTHON_H +#define LMP_MLIAP_MODEL_PYTHON_H + +#include "mliap_model.h" + +namespace LAMMPS_NS { + + +class MLIAPModelPython : public MLIAPModel { +public: + MLIAPModelPython(LAMMPS*, char* = NULL); + ~MLIAPModelPython(); + virtual int get_nparams(); + virtual int get_gamma_nnz(class MLIAPData*); + virtual void compute_gradients(class MLIAPData*); + 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: + +}; + +} + +#endif + diff --git a/src/MLIAPPY/mliap_model_python_couple.pyx b/src/MLIAPPY/mliap_model_python_couple.pyx new file mode 100644 index 0000000000..ad95ea2c48 --- /dev/null +++ b/src/MLIAPPY/mliap_model_python_couple.pyx @@ -0,0 +1,123 @@ +# 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 + +import pickle + +# For converting C arrays to numpy arrays +import numpy as np + +# For converting void * to integer for tracking object identity +from libc.stdint cimport uintptr_t + +from libcpp.string cimport string + + +cdef extern from "mliap_data.h" namespace "LAMMPS_NS": + cdef cppclass MLIAPData: + # Array shapes + int natoms + int ndescriptors + + # Input data + int * ielems # types for all atoms in list + double ** descriptors # descriptors for all atoms in list + + # Output data to write to + double ** betas # betas for all atoms in list + double * eatoms # energy for all atoms in list + double energy + +cdef extern from "mliap_model_python.h" namespace "LAMMPS_NS": + cdef cppclass MLIAPModelPython: + void connect_param_counts() + + +class MLIAPPYModelNotLinked(Exception): pass + + +LOADED_MODELS = {} + +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( 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 = c_id + lmp_model.connect_param_counts() + + +cdef public void MLIAPPY_unload_model(MLIAPModelPython * c_model) with gil: + del LOADED_MODELS[c_id(c_model)] + +cdef public int MLIAPPY_nparams(MLIAPModelPython * c_model) with gil: + return int(retrieve(c_model).n_params) + +cdef public int MLIAPPY_nelements(MLIAPModelPython * c_model) with gil: + return int(retrieve(c_model).n_elements) + +cdef public int MLIAPPY_ndescriptors(MLIAPModelPython * c_model) with gil: + return int(retrieve(c_model).n_descriptors) + +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 + + # Make numpy arrays from pointers + beta_np = np.asarray( &data.betas[0][0]) + desc_np = np.asarray( &data.descriptors[0][0]) + elem_np = np.asarray( &data.ielems[0]) + en_np = np.asarray( &data.eatoms[0]) + + # Invoke python model on numpy arrays. + model(elem_np,desc_np,beta_np,en_np) + + # Get the total energy from the atom energy. + energy = np.sum(en_np) + data.energy = energy + return diff --git a/src/MLIAPPY/mliappy_pytorch.py b/src/MLIAPPY/mliappy_pytorch.py new file mode 100644 index 0000000000..5294344c10 --- /dev/null +++ b/src/MLIAPPY/mliappy_pytorch.py @@ -0,0 +1,46 @@ +import numpy as np +import torch + +def calc_n_params(model): + return sum(p.nelement() for p in model.parameters()) + +class TorchWrapper(torch.nn.Module): + def __init__(self, model,n_descriptors,n_elements,n_params=None): + super().__init__() + self.model = model + self.model.to(self.dtype) + if n_params is None: + n_params = calc_n_params(model) + self.n_params = n_params + self.n_descriptors = n_descriptors + self.n_elements = n_elements + + def __call__(self, elems, bispectrum, beta, energy): + + bispectrum = torch.from_numpy(bispectrum).to(self.dtype).requires_grad_(True) + elems = torch.from_numpy(elems).to(torch.long) - 1 + + with torch.autograd.enable_grad(): + + energy_nn = self.model(bispectrum, elems) + if energy_nn.ndim > 1: + energy_nn = energy_nn.flatten() + + beta_nn = torch.autograd.grad(energy_nn.sum(), bispectrum)[0] + + beta[:] = beta_nn.detach().cpu().numpy().astype(np.float64) + energy[:] = energy_nn.detach().cpu().numpy().astype(np.float64) + +class TorchWrapper32(TorchWrapper): + dtype = torch.float32 + +class TorchWrapper64(TorchWrapper): + dtype = torch.float64 + +class IgnoreElems(torch.nn.Module): + def __init__(self,subnet): + super().__init__() + self.subnet = subnet + + def forward(self,bispectrum,elems): + return self.subnet(bispectrum) diff --git a/src/Makefile b/src/Makefile index 149dedd35b..8667f3f725 100644 --- a/src/Makefile +++ b/src/Makefile @@ -48,7 +48,7 @@ endif PACKAGE = asphere body class2 colloid compress coreshell dipole gpu \ granular kim kokkos kspace latte manybody mc message misc \ - mliap molecule mpiio mscg opt peri poems \ + mliap mliappy molecule mpiio mscg opt peri poems \ python qeq replica rigid shock snap spin srd voronoi PACKUSER = user-adios user-atc user-awpmd user-bocs user-cgdna user-cgsdk user-colvars \ diff --git a/src/PYTHON/python_impl.cpp b/src/PYTHON/python_impl.cpp index 22bf8e77fb..62f19d0f12 100644 --- a/src/PYTHON/python_impl.cpp +++ b/src/PYTHON/python_impl.cpp @@ -26,6 +26,14 @@ #include #include // IWYU pragma: export +#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 + using namespace LAMMPS_NS; enum{NONE,INT,DOUBLE,STRING,PTR}; @@ -47,11 +55,17 @@ 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(); - + + #ifdef LMP_MLIAPPY + // 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); + if (err) error->all(FLERR,"Could not register MLIAPPY embedded python module."); + #endif + Py_Initialize(); PyEval_InitThreads();