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..842c43bd96 --- /dev/null +++ b/cmake/Modules/Packages/MLIAPPY.cmake @@ -0,0 +1,12 @@ +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) + diff --git a/doc/src/Packages_details.rst b/doc/src/Packages_details.rst index 74214e2e0e..edce17b9ab 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,37 @@ 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, 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. + +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. + +This package includes more options for the mliap compute and pair style. + +**Author:** Nicholas Lubbers (LANL). + +**Supporting info:** + +* src/MLIAPPY: filenames -> commands +* :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.txt b/examples/mliappy/README.txt new file mode 100644 index 0000000000..f624acd657 --- /dev/null +++ b/examples/mliappy/README.txt @@ -0,0 +1,11 @@ +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 \ No newline at end of file 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..bdc3887282 --- /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.IgnoreTypes(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) \ No newline at end of file diff --git a/examples/mliappy/in.mliap.pytorch.Ta06A b/examples/mliappy/in.mliap.pytorch.Ta06A new file mode 100644 index 0000000000..526c24ab23 --- /dev/null +++ b/examples/mliappy/in.mliap.pytorch.Ta06A @@ -0,0 +1,67 @@ +# Demonstrate MLIAP interface to kinear 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 + +python simple here """ +from __future__ import print_function + +def simple(): + foo = 0 + print("Inside simple function") + try: + foo += 1 + except Exception as e: + print("FOO error:", e) +""" + +python simple invoke + +# 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/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_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 d3c8373708..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; } 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 36642c1a87..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; } 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 ef3ff40da5..9fd7b483af 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; @@ -109,6 +113,7 @@ void PairMLIAP::allocate() void PairMLIAP::settings(int narg, char ** arg) { + if (narg < 4) error->all(FLERR,"Illegal pair_style command"); @@ -132,6 +137,13 @@ 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) { @@ -215,10 +227,13 @@ void PairMLIAP::coeff(int narg, char **arg) // 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"); + 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)"); + }; } /* ---------------------------------------------------------------------- diff --git a/src/MLIAPPY/mliap_model_python.cpp b/src/MLIAPPY/mliap_model_python.cpp new file mode 100644 index 0000000000..13e5a85e83 --- /dev/null +++ b/src/MLIAPPY/mliap_model_python.cpp @@ -0,0 +1,163 @@ +/* ---------------------------------------------------------------------- + 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) +{ + int err; + + 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"); + 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() +{ + 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) { + PyErr_Print(); + PyErr_Clear(); + PyGILState_Release(gstate); + error->all(FLERR,"Loading python model failure."); + } + nelements = MLIAPPY_nelements(this); + nparams = MLIAPPY_nparams(this); + ndescriptors = MLIAPPY_ndescriptors(this); +} + +/* ---------------------------------------------------------------------- + Calculate model gradients w.r.t descriptors + for each atom beta_i = dE(B_i)/dB_i + ---------------------------------------------------------------------- */ + +void MLIAPModelPython::compute_gradients(MLIAPData* data) +{ + MLIAPPY_model_callback(this, data); +} + +/* ---------------------------------------------------------------------- + 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..3b162509f7 --- /dev/null +++ b/src/MLIAPPY/mliap_model_python.h @@ -0,0 +1,42 @@ +/* -*- 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(); +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..3e03247fa7 --- /dev/null +++ b/src/MLIAPPY/mliap_model_python_couple.pyx @@ -0,0 +1,85 @@ +# cython: language_level=3 +# distutils: language = c++ + +cimport cython + +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 + +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: + ctypedef void (*CBPtr)(void * , MLIAPData); + void set_model(CBPtr, void *); + + +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( c_model)] = model + return 0 + +cdef public void MLIAPPY_unload_model(MLIAPModelPython * c_model) with gil: + del LOADED_MODELS[int( c_model)] + +cdef public int MLIAPPY_nparams(MLIAPModelPython * c_model) with gil: + model = LOADED_MODELS[int( c_model)] + n_params = int(model.n_params) + return n_params + +cdef public int MLIAPPY_nelements(MLIAPModelPython * c_model) with gil: + model = LOADED_MODELS[int( c_model)] + n_elements = int(model.n_elements) + return n_elements + +cdef public int MLIAPPY_ndescriptors(MLIAPModelPython * c_model) with gil: + model = LOADED_MODELS[int( c_model)] + n_descriptors = int(model.n_descriptors) + return n_descriptors + +cdef public MLIAPPY_model_callback(MLIAPModelPython * c_model, MLIAPData * data) with gil: + model = LOADED_MODELS[int( 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]) + type_np = np.asarray( &data.ielems[0]) + en_np = np.asarray( &data.eatoms[0]) + + # Invoke python model on numpy arrays. + model(type_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..b9f7ac3dcc --- /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, types, bispectrum, beta, energy): + + bispectrum = torch.from_numpy(bispectrum).to(self.dtype).requires_grad_(True) + types = torch.from_numpy(types).to(torch.long) - 1 + + with torch.autograd.enable_grad(): + + energy_nn = self.model(bispectrum, types) + 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 IgnoreTypes(torch.nn.Module): + def __init__(self,subnet): + super().__init__() + self.subnet = subnet + + def forward(self,bispectrum,types): + 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..9f54822c47 100644 --- a/src/PYTHON/python_impl.cpp +++ b/src/PYTHON/python_impl.cpp @@ -26,6 +26,11 @@ #include #include // IWYU pragma: export +#ifdef LMP_MLIAPPY +#include "mliap_model_python.h" +#include "mliap_model_python_couple.h" +#endif + using namespace LAMMPS_NS; enum{NONE,INT,DOUBLE,STRING,PTR}; @@ -51,7 +56,14 @@ PythonImpl::PythonImpl(LAMMPS *lmp) : Pointers(lmp) // 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); + // todo: catch if error and report problem. + #endif + Py_Initialize(); PyEval_InitThreads();