Merge remote-tracking branch 'origin/mliappy2' into mliappy3

.gitignore

@@ -46,3 +46,4 @@ Thumbs.db
 /Makefile
 /cmake_install.cmake
 /lmp
+

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})

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)

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)

doc/src/Packages_details.rst

@@ -45,6 +45,7 @@ page gives those details.
    * :ref:`MESSAGE <PKG-MESSAGE>`
    * :ref:`MISC <PKG-MISC>`
    * :ref:`MLIAP <PKG-MLIAP>`
+   * :ref:`MLIAPPY <PKG-MLIAPPY>`
    * :ref:`MOLECULE <PKG-MOLECULE>`
    * :ref:`MPIIO <PKG-MPIIO>`
    * :ref:`MSCG <PKG-MSCG>`
@@ -678,6 +679,36 @@ To use this package, also the :ref:`SNAP package <PKG-SNAP>` 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 <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.
+
+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 <pair_mliap>`
+* examples/mliappy (see README)
+----------
+
 .. _PKG-MOLECULE:
 
 MOLECULE package

doc/src/Packages_standard.rst

@@ -61,6 +61,8 @@ package:
 +----------------------------------+--------------------------------------+----------------------------------------------------+------------------------------------------------------+---------+
 | :ref:`MLIAP <PKG-MLIAP>`         | multiple machine learning potentials | :doc:`pair_style mliap <pair_mliap>`               | mliap                                                | no      |
 +----------------------------------+--------------------------------------+----------------------------------------------------+------------------------------------------------------+---------+
+| :ref:`MLIAPPY <PKG-MLIAPPY>`     | enable python ML models for mliap    | :doc:`pair_style mliap <pair_mliap>`               | mliap                                                | no      |
++----------------------------------+--------------------------------------+----------------------------------------------------+------------------------------------------------------+---------+
 | :ref:`MOLECULE <PKG-MOLECULE>`   | molecular system force fields        | :doc:`Howto bioFF <Howto_bioFF>`                   | peptide                                              | no      |
 +----------------------------------+--------------------------------------+----------------------------------------------------+------------------------------------------------------+---------+
 | :ref:`MPIIO <PKG-MPIIO>`         | MPI parallel I/O dump and restart    | :doc:`dump <dump>`                                 | n/a                                                  | no      |

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 <Build_package>` doc page for more info.
 
+Python models such as neural networks can be used if the MLIAPPY package is built.
+
 Related commands
 """"""""""""""""
 

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 <pair_snap>`, including the linear, quadratic,
-and chem variants. Work is currently underway to extend
+and chem variants. With the MLIAPPY package installed, the *mliappy* model 
-the interface to handle neural network energy models,
+is available which can be used to couple python models such as neural network
-and it is also straightforward to add new descriptor styles.
+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 <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 <Build_package>` doc page for more info.
 
+Python models such as neural networks can be used if the MLIAPPY package is built.
+
+
 Related commands
 """"""""""""""""
 

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`

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
+

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
+

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)

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}
+

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)
+
+

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")

python/lammps.py

@@ -506,6 +506,8 @@ class lammps(object):
     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)
+            

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");

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

src/MLIAP/mliap_data.h

@@ -36,6 +36,8 @@ class MLIAPData : protected Pointers {
   double **gradforce;
   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

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;
-  }
   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;
 

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 *);
+
 };
 
 }

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;
     }
   }
 }

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();

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;
     }
   }
 }

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();

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 <cmath>
 #include <cstring>
+#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_global) eng_vdwl += data->energy;
-  if (eflag_atom) eatom[i] += ei;
+  if (eflag_atom) 
+    for (int ii = 0; ii < data->natoms; ii++) {
+      const int i = data->iatoms[ii];
+      eatom[i] += data->eatoms[ii];
+    }
 }
 
 /* ----------------------------------------------------------------------

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);

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

src/MLIAPPY/Makefile.lammps

@@ -0,0 +1 @@
+#TODO

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 <Python.h>
+#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;
+}

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
+

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(<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[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(<double[:n_a,:n_d] > &data.betas[0][0])
+    desc_np = np.asarray(<double[:n_a,:n_d]> &data.descriptors[0][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(elem_np,desc_np,beta_np,en_np)
+    
+    # Get the total energy from the atom energy.
+    energy = np.sum(en_np)
+    data.energy = <double> energy
+    return

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)

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 \

src/PYTHON/python_impl.cpp

@@ -26,6 +26,14 @@
 #include <cstring>
 #include <Python.h>  // 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();