Initial commit of mliap unified work

cmake/Modules/Packages/ML-IAP.cmake

@@ -25,16 +25,18 @@ if(MLIAP_ENABLE_PYTHON)
   endif()
 
   set(MLIAP_BINARY_DIR ${CMAKE_BINARY_DIR}/cython)
-  set(MLIAP_CYTHON_SRC ${LAMMPS_SOURCE_DIR}/ML-IAP/mliap_model_python_couple.pyx)
+  file(GLOB MLIAP_CYTHON_SRC ${LAMMPS_SOURCE_DIR}/ML-IAP/*.pyx)
-  get_filename_component(MLIAP_CYTHON_BASE ${MLIAP_CYTHON_SRC} NAME_WE)
   file(MAKE_DIRECTORY ${MLIAP_BINARY_DIR})
-  add_custom_command(OUTPUT  ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.cpp ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.h
+  foreach(MLIAP_CYTHON_FILE ${MLIAP_CYTHON_SRC})
-          COMMAND            ${CMAKE_COMMAND} -E copy_if_different ${MLIAP_CYTHON_SRC} ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.pyx
+    get_filename_component(MLIAP_CYTHON_BASE ${MLIAP_CYTHON_FILE} NAME_WE)
-          COMMAND            ${Cythonize_EXECUTABLE} -3 ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.pyx
+    add_custom_command(OUTPUT  ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.cpp ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.h
-          WORKING_DIRECTORY  ${MLIAP_BINARY_DIR}
+            COMMAND            ${CMAKE_COMMAND} -E copy_if_different ${MLIAP_CYTHON_FILE} ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.pyx
-          MAIN_DEPENDENCY    ${MLIAP_CYTHON_SRC}
+            COMMAND            ${Cythonize_EXECUTABLE} -3 ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.pyx
-          COMMENT "Generating C++ sources with cythonize...")
+            WORKING_DIRECTORY  ${MLIAP_BINARY_DIR}
+            MAIN_DEPENDENCY    ${MLIAP_CYTHON_FILE}
+            COMMENT "Generating C++ sources with cythonize...")
+    target_sources(lammps PRIVATE ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.cpp)
+  endforeach()
   target_compile_definitions(lammps PRIVATE -DMLIAP_PYTHON)
-  target_sources(lammps PRIVATE ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.cpp)
   target_include_directories(lammps PRIVATE ${MLIAP_BINARY_DIR})
 endif()

examples/mliap/InP.hippynn.mliap.unified

@@ -0,0 +1,5 @@
+# Definition of HIPNN potential
+ 
+# Specify MLIAP Unified pair style
+pair_style mliap unified mliap_unified_hippynn_InP.pkl
+pair_coeff * * In P

examples/mliap/in.mliap.unified.hippynn.Al

@@ -0,0 +1,56 @@
+# Demonstrate MLIAP interface to HIPNN Al potential
+
+# Initialize simulation
+
+variable nsteps index 100
+variable nrep equal 4
+variable a equal 4.05
+units           metal
+
+# generate the box and atom positions using a FCC lattice
+
+variable nx equal ${nrep}
+variable ny equal ${nrep}
+variable nz equal ${nrep}
+
+boundary        p p p
+
+lattice         fcc $a
+region          box block 0 ${nx} 0 ${ny} 0 ${nz}
+create_box      1 box
+create_atoms    1 box
+
+mass 1 26.981
+
+# choose potential
+
+pair_style	mliap unified mliap_unified_hippynn_Al.pkl
+pair_coeff	* * Al
+
+# 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
+
+#dump 4 all custom 1 forces.xyz fx fy fz
+
+# Run MD
+
+velocity all create 300.0 4928459 loop geom
+fix 1 all nve
+run             ${nsteps}
+

examples/mliap/in.mliap.unified.hippynn.InP

@@ -0,0 +1,52 @@
+# Demonstrate MLIAP interface to HIPNN InP potential
+
+# Initialize simulation
+
+variable nsteps index 100
+variable nrep equal 4
+variable a equal 5.83
+units           metal
+
+# generate the box and atom positions using a FCC lattice
+
+variable nx equal ${nrep}
+variable ny equal ${nrep}
+variable nz equal ${nrep}
+
+boundary        p p p
+
+lattice         diamond $a
+region          box block 0 ${nx} 0 ${ny} 0 ${nz}
+create_box      2 box
+create_atoms    1 box basis 5 2 basis 6 2 basis 7 2 basis 8 2
+
+mass 1 114.76
+mass 2 30.98
+
+# choose potential
+
+include InP.hippynn.mliap.unified
+
+#dump 4 all custom 1 forces.xyz fx fy fz
+
+#dump		3 all movie 1 movie.mpg type type &
+#		axes yes 0.8 0.02 view 60 -30
+#dump_modify	3 pad 3
+
+
+# Setup output
+
+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/mliap/in.mliap.unified.lj.Ar

@@ -0,0 +1,35 @@
+# 3d Lennard-Jones melt
+
+units		lj
+atom_style	atomic
+
+lattice		fcc 0.8442
+region		box block 0 10 0 10 0 10
+create_box	1 box
+create_atoms	1 box
+mass		1 1.0
+
+velocity	all create 3.0 87287 loop geom
+
+pair_style	mliap unified mliap_unified_lj_Ar.pkl
+pair_coeff	* * Ar 2.5
+
+neighbor	0.3 bin
+neigh_modify	every 1 delay 0 check no
+
+fix		1 all nve
+
+#dump		id all atom 50 dump.melt
+
+#dump		2 all image 25 image.*.jpg type type &
+#		axes yes 0.8 0.02 view 60 -30
+#dump_modify	2 pad 3
+
+#dump		3 all movie 1 movie.mpg type type &
+#		axes yes 0.8 0.02 view 60 -30
+#dump_modify	3 pad 3
+
+#dump 4 all custom 1 forces.xyz fx fy fz
+
+thermo		100
+run		25000

examples/mliap/pickle_mliap_unified_lj_Al.py

@@ -0,0 +1,21 @@
+import pickle
+import lammps
+import lammps.mliap
+
+from lammps.mliap.mliap_unified_lj import MLIAPUnifiedLJ
+
+
+if __name__ == '__main__':
+    unified = MLIAPUnifiedLJ()
+    unified.element_types = ["Al"]
+    unified.ndescriptors = 1
+    unified.nparams = 3
+    # Mimicking the LJ pair-style:
+    # pair_style lj/cut 2.5
+    # pair_coeff * * 1 1
+    unified.epsilon = 1.0
+    unified.sigma = 2.5
+    unified.rcutfac = 5.0
+
+    with open('mliap_unified_lj_Al.pkl', 'wb') as fp:
+        pickle.dump(unified, fp)

examples/mliap/pickle_mliap_unified_lj_Ar.py

@@ -0,0 +1,21 @@
+import pickle
+import lammps
+import lammps.mliap
+
+from lammps.mliap.mliap_unified_lj import MLIAPUnifiedLJ
+
+
+if __name__ == '__main__':
+    unified = MLIAPUnifiedLJ()
+    unified.element_types = ["Ar"]
+    unified.ndescriptors = 1
+    unified.nparams = 3
+    # Mimicking the LJ pair-style:
+    # pair_style lj/cut 2.5
+    # pair_coeff * * 1 1
+    unified.epsilon = 1.0
+    unified.sigma = 1.0
+    unified.rcutfac = 1.25
+
+    with open('mliap_unified_lj_Ar.pkl', 'wb') as fp:
+        pickle.dump(unified, fp)

python/lammps/mliap/loader.py

@@ -17,27 +17,58 @@
 
 
 import sys
+import importlib
 import importlib.util
 import importlib.machinery
+import importlib.abc
+
+from ctypes import pythonapi, c_int, c_void_p, py_object
+
+# This dynamic loader imports a python module embedded in a shared library.
+# The default value of api_version is 1013 because it has been stable since 2006.
+class DynamicLoader(importlib.abc.Loader):
+    def __init__(self,module_name,library,api_version=1013):
+        self.api_version = api_version
+        
+        attr = "PyInit_"+module_name
+        initfunc = getattr(library,attr)
+        # c_void_p is standin for PyModuleDef *
+        initfunc.restype = c_void_p
+        initfunc.argtypes = ()
+        self.module_def = initfunc()
+
+    def create_module(self, spec):
+        createfunc = pythonapi.PyModule_FromDefAndSpec2
+        # c_void_p is standin for PyModuleDef *
+        createfunc.argtypes = c_void_p, py_object, c_int
+        createfunc.restype = py_object
+        module = createfunc(self.module_def, spec, self.api_version)
+        return module
+    
+    def exec_module(self, module):
+        execfunc = pythonapi.PyModule_ExecDef
+        # c_void_p is standin for PyModuleDef *
+        execfunc.argtypes = py_object, c_void_p
+        execfunc.restype = c_int
+        result = execfunc(module, self.module_def)
+        if result<0:
+           raise ImportError()
 
 def activate_mliappy(lmp):
     try:
-        # Begin Importlib magic to find the embedded python module
+        library = lmp.lib
-        # This is needed because the filename for liblammps does not
+        module_names = ["mliap_model_python_couple", "mliap_unifiedpy"]
-        # match the spec for normal python modules, wherein
+        api_version = library.lammps_PYTHON_API_VERSION()
-        # file names match with PyInit function names.
+        
-        # Also, python normally doesn't look for extensions besides '.so'
+        for module_name in module_names:
-        # We fix both of these problems by providing an explict
+            # Make Machinery
-        # path to the extension module 'mliap_model_python_couple' in
+            loader = DynamicLoader(module_name,library,api_version)
-
+            spec = importlib.util.spec_from_loader(module_name,loader)
-        path = lmp.lib._name
+            
-        loader = importlib.machinery.ExtensionFileLoader('mliap_model_python_couple', path)
+            # Do the import
-        spec = importlib.util.spec_from_loader('mliap_model_python_couple', loader)
+            module = importlib.util.module_from_spec(spec)
-        module = importlib.util.module_from_spec(spec)
+            sys.modules[module_name] = module
-        sys.modules['mliap_model_python_couple'] = module
+            spec.loader.exec_module(module)
-        spec.loader.exec_module(module)
-        # End Importlib magic to find the embedded python module
-
     except Exception as ee:
         raise ImportError("Could not load ML-IAP python coupling module.") from ee
 
@@ -49,4 +80,3 @@ def load_model(model):
                           "the pair style. Call lammps.mliap.activate_mliappy(lmp)."
                           ) from ie
     mliap_model_python_couple.load_from_python(model)
-

python/lammps/mliap/mliap_unified_abc.py

@@ -0,0 +1,23 @@
+from abc import ABC, abstractmethod
+
+class MLIAPUnified(ABC):
+    """Abstract base class for MLIAPUnified."""
+
+    def __init__(self):
+        self.interface = None
+        self.element_types = None
+        self.ndescriptors = None
+        self.nparams = None
+        self.rcutfac = None
+
+    @abstractmethod
+    def compute_gradients(self, data):
+        """Compute gradients."""
+
+    @abstractmethod
+    def compute_descriptors(self, data):
+        """Compute descriptors."""
+
+    @abstractmethod
+    def compute_forces(self, data):
+        """Compute forces."""

python/lammps/mliap/mliap_unified_lj.py

@@ -0,0 +1,35 @@
+from .mliap_unified_abc import MLIAPUnified
+import numpy as np
+
+
+class MLIAPUnifiedLJ(MLIAPUnified):
+    """Test implementation for MLIAPUnified."""
+
+    def __init__(self):
+        super().__init__()
+
+    def compute_gradients(self, data):
+        """Test compute_gradients."""
+    
+    def compute_descriptors(self, data):
+        """Test compute_descriptors."""
+    
+    def compute_forces(self, data):
+        """Test compute_forces."""
+        eij, fij = self.compute_pair_ef(data)
+        data.update_pair_energy(eij)
+        data.update_pair_forces(fij)
+ 
+    def compute_pair_ef(self, data):
+        rij = data.rij
+
+        r2inv = 1.0 / np.sum(rij ** 2, axis=1)
+        r6inv = r2inv * r2inv * r2inv
+
+        lj1 = 4.0 * self.epsilon * self.sigma**12
+        lj2 = 4.0 * self.epsilon * self.sigma**6
+
+        eij = r6inv * (lj1 * r6inv - lj2)
+        fij = r6inv * (3.0 * lj2 - 6.0 * lj2 * r6inv) * r2inv
+        fij = fij[:, np.newaxis] * rij
+        return eij, fij

src/ML-IAP/mliap_data.cpp

@@ -31,12 +31,14 @@ 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),
+  Pointers(lmp), f(nullptr), 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),
+  numneighs(nullptr), iatoms(nullptr), pair_i(nullptr), ielems(nullptr),
-  rij(nullptr), graddesc(nullptr), model(nullptr), descriptor(nullptr), list(nullptr)
+  jatoms(nullptr), jelems(nullptr), elems(nullptr), rij(nullptr),
+  graddesc(nullptr), model(nullptr), descriptor(nullptr), list(nullptr)
 {
+  f = atom->f;
   gradgradflag = gradgradflag_in;
   map = map_in;
   model = model_in;
@@ -85,10 +87,12 @@ MLIAPData::~MLIAPData()
   memory->destroy(gradforce);
 
   memory->destroy(iatoms);
+  memory->destroy(pair_i);
   memory->destroy(ielems);
   memory->destroy(numneighs);
   memory->destroy(jatoms);
   memory->destroy(jelems);
+  memory->destroy(elems);
   memory->destroy(rij);
   memory->destroy(graddesc);
 }
@@ -110,25 +114,32 @@ void MLIAPData::generate_neighdata(NeighList* list_in, int eflag_in, int vflag_i
   double **x = atom->x;
   int *type = atom->type;
 
+  nlocal = atom->nlocal;
+  nghost = atom->nghost;
+  ntotal = nlocal + nghost;
+
   int *ilist = list->ilist;
   int *numneigh = list->numneigh;
   int **firstneigh = list->firstneigh;
 
-  // grow nmax gradforce array if necessary
+  // grow nmax gradforce, elems arrays if necessary
 
   if (atom->nmax > nmax) {
     nmax = atom->nmax;
     memory->grow(gradforce,nmax,size_gradforce,
                  "MLIAPData:gradforce");
+    memory->grow(elems,nmax,"MLIAPData:elems");
  }
 
-  // clear gradforce array
+  // clear gradforce, elems arrays
 
   int nall = atom->nlocal + atom->nghost;
-  for (int i = 0; i < nall; i++)
+  for (int i = 0; i < nall; i++) {
     for (int j = 0; j < size_gradforce; j++) {
       gradforce[i][j] = 0.0;
     }
+    elems[i] = 0;
+  }
 
   // grow arrays if necessary
 
@@ -153,6 +164,7 @@ void MLIAPData::generate_neighdata(NeighList* list_in, int eflag_in, int vflag_i
 
   grow_neigharrays();
 
+  npairs = 0;
   int ij = 0;
   for (int ii = 0; ii < nlistatoms; ii++) {
     const int i = ilist[ii];
@@ -178,6 +190,7 @@ void MLIAPData::generate_neighdata(NeighList* list_in, int eflag_in, int vflag_i
       const int jelem = map[jtype];
 
       if (rsq < descriptor->cutsq[ielem][jelem]) {
+        pair_i[ij] = i;
         jatoms[ij] = j;
         jelems[ij] = jelem;
         rij[ij][0] = delx;
@@ -190,6 +203,12 @@ void MLIAPData::generate_neighdata(NeighList* list_in, int eflag_in, int vflag_i
     iatoms[ii] = i;
     ielems[ii] = ielem;
     numneighs[ii] = ninside;
+    npairs += ninside;
+  }
+
+  for (int i = 0; i < nall; i++) {
+    const int itype = type[i];
+    elems[i] = map[itype];
   }
 
   eflag = eflag_in;
@@ -249,6 +268,7 @@ void MLIAPData::grow_neigharrays()
   }
 
   if (nneigh_max < nneigh) {
+    memory->grow(pair_i,nneigh,"MLIAPData:pair_i");
     memory->grow(jatoms,nneigh,"MLIAPData:jatoms");
     memory->grow(jelems,nneigh,"MLIAPData:jelems");
     memory->grow(rij,nneigh,3,"MLIAPData:rij");
@@ -264,6 +284,7 @@ double MLIAPData::memory_usage()
 
   bytes += (double)nelements*nparams*sizeof(double);     // egradient
   bytes += (double)nmax*size_gradforce*sizeof(double);   // gradforce
+  bytes += (double)nmax*sizeof(int);                     // elems
 
   if (gradgradflag == 1) {
     bytes += (double)natomgamma_max*
@@ -282,6 +303,7 @@ double MLIAPData::memory_usage()
   bytes += (double)natomneigh_max*sizeof(int);               // ielems
   bytes += (double)natomneigh_max*sizeof(int);               // numneighs
 
+  bytes += (double)nneigh_max*sizeof(int);                   // pair_i
   bytes += (double)nneigh_max*sizeof(int);                   // jatoms
   bytes += (double)nneigh_max*sizeof(int);                   // jelems
   bytes += (double)nneigh_max*3*sizeof(double);              // rij"

src/ML-IAP/mliap_data.h

@@ -35,6 +35,7 @@ class MLIAPData : protected Pointers {
   int size_gradforce;
   int yoffset, zoffset;
   int ndims_force, ndims_virial;
+  double **f;
   double **gradforce;
   double **betas;          // betas for all atoms in list
   double **descriptors;    // descriptors for all atoms in list
@@ -56,15 +57,19 @@ class MLIAPData : protected Pointers {
   // data structures for mliap neighbor list
   // only neighbors strictly inside descriptor cutoff
 
+  int nlocal, nghost, ntotal;    // # of owned and ghost atoms on this proc, and sum of both
   int nlistatoms;                // current number of atoms in neighborlist
   int nlistatoms_max;            // allocated size of descriptor array
   int natomneigh_max;            // allocated size of atom neighbor arrays
   int *numneighs;                // neighbors count for each atom
   int *iatoms;                   // index of each atom
+  int *pair_i;                   // index of each i atom for each ij pair
   int *ielems;                   // element of each atom
   int nneigh_max;                // number of ij neighbors allocated
+  int npairs;                    // number of ij neighbor pairs
   int *jatoms;                   // index of each neighbor
   int *jelems;                   // element of each neighbor
+  int *elems;                    // element of each atom in or not in the neighborlist
   double **rij;                  // distance vector of each neighbor
   double ***graddesc;            // descriptor gradient w.r.t. each neighbor
   int eflag;                     // indicates if energy is needed

src/ML-IAP/mliap_unified.cpp

@@ -0,0 +1,234 @@
+/* ----------------------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+   https://www.lammps.org/, 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.
+------------------------------------------------------------------------- */
+
+/* ----------------------------------------------------------------------
+   Contributing author: Steven Ray Anaya (LANL)
+------------------------------------------------------------------------- */
+#ifdef MLIAP_PYTHON
+
+#include <Python.h>
+#include "mliap_unified.h"
+
+#include "error.h"
+#include "memory.h"
+#include "lmppython.h"
+#include "mliap_data.h"
+#include "mliap_unifiedpy.h"
+#include "pair_mliap.h"
+#include "python_compat.h"
+#include "utils.h"
+
+using namespace LAMMPS_NS;
+
+/* ---------------------------------------------------------------------- */
+
+MLIAPDummyDescriptor::MLIAPDummyDescriptor(LAMMPS *lmp) :
+    MLIAPDescriptor(lmp) {}
+
+MLIAPDummyDescriptor::~MLIAPDummyDescriptor() {
+  memory->destroy(radelem);
+  memory->destroy(cutsq);
+  Py_DECREF(unified_interface);
+}
+
+void MLIAPDummyDescriptor::compute_descriptors(class MLIAPData *data) {
+  PyGILState_STATE gstate = PyGILState_Ensure();
+  compute_descriptors_python(unified_interface, data);
+  if (PyErr_Occurred()) {
+    PyErr_Print();
+    PyErr_Clear();
+    PyGILState_Release(gstate);
+    lmp->error->all(FLERR, "Running mliappy unified compute_descriptors failure.");
+  }
+  PyGILState_Release(gstate);
+}
+
+void MLIAPDummyDescriptor::compute_forces(class MLIAPData *data) {
+  PyGILState_STATE gstate = PyGILState_Ensure();
+  compute_forces_python(unified_interface, data);
+  if (PyErr_Occurred()) {
+    PyErr_Print();
+    PyErr_Clear();
+    PyGILState_Release(gstate);
+    lmp->error->all(FLERR, "Running mliappy unified compute_forces failure.");
+  }
+  PyGILState_Release(gstate);
+}
+
+void MLIAPDummyDescriptor::compute_force_gradients(class MLIAPData *) {
+  error->all(FLERR, "compute_force_gradients not implemented");
+}
+
+void MLIAPDummyDescriptor::compute_descriptor_gradients(class MLIAPData *) {
+  error->all(FLERR, "compute_descriptor_gradients not implemented");
+}
+
+void MLIAPDummyDescriptor::init() {
+  memory->create(radelem, nelements, "mliap_dummy_descriptor:radelem");
+  for (int ielem = 0; ielem < nelements; ielem++) {
+    radelem[ielem] = 1;
+  }
+
+  double cut;
+  cutmax = 0.0;
+  memory->create(cutsq, nelements, nelements, "mliap/descriptor/dummy:cutsq");
+  for (int ielem = 0; ielem < nelements; ielem++) {
+    // rcutfac set from python, is global cutoff for all elements
+    cut = 2.0*radelem[ielem]*rcutfac;
+    if (cut > cutmax) cutmax = cut;
+    cutsq[ielem][ielem] = cut*cut;
+    for (int jelem = ielem+1; jelem < nelements; jelem++) {
+      cut = (radelem[ielem]+radelem[jelem])*rcutfac;
+      cutsq[ielem][jelem] = cutsq[jelem][ielem] = cut*cut;
+    }
+  }
+}
+
+void MLIAPDummyDescriptor::set_elements(char **elems, int nelems) {
+  nelements = nelems;
+  elements = new char*[nelems];
+  for (int i = 0; i < nelems; i++) {
+    elements[i] = utils::strdup(elems[i]);
+  }
+}
+
+/* ---------------------------------------------------------------------- */
+
+MLIAPDummyModel::MLIAPDummyModel(LAMMPS *lmp, char *coefffilename) :
+    MLIAPModel(lmp, coefffilename) {
+  nonlinearflag = 1;
+}
+
+MLIAPDummyModel::~MLIAPDummyModel() {
+  // MLIAPPY_unload_model(this);
+  Py_DECREF(unified_interface);
+}
+
+int MLIAPDummyModel::get_nparams() {
+  return nparams;
+}
+
+int MLIAPDummyModel::get_gamma_nnz(class MLIAPData *) {
+  // TODO: get_gamma_nnz
+  return 0;
+}
+
+void MLIAPDummyModel::compute_gradients(class MLIAPData *data) {
+  PyGILState_STATE gstate = PyGILState_Ensure();
+  compute_gradients_python(unified_interface, data);
+  if (PyErr_Occurred()) {
+    PyErr_Print();
+    PyErr_Clear();
+    PyGILState_Release(gstate);
+    lmp->error->all(FLERR, "Running mliappy unified compute_gradients failure.");
+  }
+  PyGILState_Release(gstate);
+}
+
+void MLIAPDummyModel::compute_gradgrads(class MLIAPData *) {
+  error->all(FLERR, "compute_gradgrads not implemented");
+}
+
+void MLIAPDummyModel::compute_force_gradients(class MLIAPData *) {
+  error->all(FLERR, "compute_force_gradients not implemented");
+}
+
+double MLIAPDummyModel::memory_usage() {
+  // TODO: implement memory usage in Cython(?)
+  return 0;
+}
+
+void MLIAPDummyModel::read_coeffs(char *) {
+  error->all(FLERR, "read_coeffs not implemented");
+}
+
+/* ---------------------------------------------------------------------- */
+
+MLIAPBuildUnified_t LAMMPS_NS::build_unified(char *unified_fname, MLIAPData *data, LAMMPS *lmp,
+    char *coefffilename) {
+  lmp->python->init();
+  PyGILState_STATE gstate = PyGILState_Ensure();
+
+  PyObject *pyMain = PyImport_AddModule("__main__");
+
+  if (!pyMain) {
+    PyGILState_Release(gstate);
+    lmp->error->all(FLERR, "Could not initialize embedded Python");
+  }
+
+  PyObject *unified_module = PyImport_ImportModule("mliap_unifiedpy");
+
+  if (!unified_module) {
+    PyErr_Print();
+    PyErr_Clear();
+    PyGILState_Release(gstate);
+    lmp->error->all(FLERR, "Loading mliappy unified module failure.");
+  }
+
+  // Connect dummy model, dummy descriptor, data to Python unified
+  MLIAPDummyModel *model = new MLIAPDummyModel(lmp, coefffilename);
+  MLIAPDummyDescriptor *descriptor = new MLIAPDummyDescriptor(lmp);
+
+  PyObject *unified_interface = mliap_unified_connect(unified_fname, model, descriptor);
+  if (PyErr_Occurred()) {
+    PyErr_Print();
+    PyErr_Clear();
+    PyGILState_Release(gstate);
+    lmp->error->all(FLERR, "Running mliappy unified module failure.");
+  }
+
+  model->unified_interface = unified_interface;
+  Py_INCREF(unified_interface);
+  descriptor->unified_interface = unified_interface;
+  Py_INCREF(unified_interface);
+
+  PyGILState_Release(gstate);
+
+  MLIAPBuildUnified_t build = {data, descriptor, model};
+  return build;
+}
+
+/* ---------------------------------------------------------------------- */
+
+void LAMMPS_NS::update_pair_energy(MLIAPData *data, double *eij) {
+    double e_total = 0;
+    for (int ii = 0; ii < data->nlistatoms; ii++)
+        data->eatoms[ii] = 0;
+
+    for (int ii = 0; ii < data->npairs; ii++) {
+        int i = data->pair_i[ii];
+        double e = 0.5 * eij[ii];
+
+        data->eatoms[i] += e;
+        e_total += e;
+    }
+    data->energy = e_total;
+}
+
+void LAMMPS_NS::update_pair_forces(MLIAPData *data, double *fij) {
+    double **f = data->f;
+    for (int ii = 0; ii < data->npairs; ii++) {
+        int ii3 = ii * 3;
+        int i = data->pair_i[ii];
+        int j = data->jatoms[ii];
+
+        f[i][0] += fij[ii3];
+        f[i][1] += fij[ii3 + 1];
+        f[i][2] += fij[ii3 + 2];
+        f[j][0] -= fij[ii3];
+        f[j][1] -= fij[ii3 + 1];
+        f[j][2] -= fij[ii3 + 2];
+    }
+}
+
+#endif

src/ML-IAP/mliap_unified.h

@@ -0,0 +1,69 @@
+/* -*- c++ -*- ----------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+   https://www.lammps.org/, 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_UNIFIED_H
+#define LMP_MLIAP_UNIFIED_H
+
+#include "mliap_data.h"
+#include "mliap_descriptor.h"
+#include "mliap_model.h"
+
+#include <Python.h>
+
+namespace LAMMPS_NS {
+
+class MLIAPDummyDescriptor : public MLIAPDescriptor {
+ public:
+  MLIAPDummyDescriptor(LAMMPS *);
+  ~MLIAPDummyDescriptor();
+  virtual void compute_descriptors(class MLIAPData *);
+  virtual void compute_forces(class MLIAPData *);
+  virtual void compute_force_gradients(class MLIAPData *);
+  virtual void compute_descriptor_gradients(class MLIAPData *);
+  virtual void init();
+  void set_elements(char **, int);
+
+  PyObject *unified_interface;
+  double rcutfac;
+};
+
+class MLIAPDummyModel : public MLIAPModel {
+ public:
+  MLIAPDummyModel(LAMMPS *, char * = NULL);
+  ~MLIAPDummyModel();
+  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();
+
+  PyObject *unified_interface;
+
+ protected:
+  virtual void read_coeffs(char *);
+};
+
+struct MLIAPBuildUnified_t {
+  MLIAPData *data;
+  MLIAPDummyDescriptor *descriptor;
+  MLIAPDummyModel *model;
+};
+
+MLIAPBuildUnified_t build_unified(char *, MLIAPData *, LAMMPS *, char * = NULL);
+void update_pair_energy(MLIAPData *, double *);
+void update_pair_forces(MLIAPData *, double *);
+
+}    // namespace LAMMPS_NS
+
+#endif

src/ML-IAP/mliap_unifiedpy.pyx

@@ -0,0 +1,375 @@
+# cython: language_level=3
+# distutils: language = c++
+
+import pickle
+import numpy as np
+import lammps.mliap
+
+cimport cython
+from cpython.ref cimport PyObject
+from libc.stdlib cimport malloc, free
+
+
+cdef extern from "../lammps.h" namespace "LAMMPS_NS":
+    cdef cppclass LAMMPS:
+        pass
+
+
+cdef extern from "mliap_data.h" namespace "LAMMPS_NS":
+    cdef cppclass MLIAPData:
+        # ----- may not need -----
+        int size_array_rows
+        int size_array_cols
+        int natoms
+        int yoffset
+        int zoffset
+        int ndims_force
+        int ndims_virial
+        # -END- may not need -END-
+        int size_gradforce
+        # ----- write only -----
+        double ** f
+        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
+        # -END- write only -END-
+        int ndescriptors        # number of descriptors
+        int nparams             # number of model parameters per element
+        int nelements           # number of elements
+
+        # data structures for grad-grad list (gamma)
+
+        # ----- ignore for now -----
+        int gamma_nnz           # number of non-zero entries in gamma
+        double ** gamma         # gamma element
+        int ** gamma_row_index  # row (parameter) index
+        int ** gamma_col_index  # column (descriptor) index
+        double * egradient      # energy gradient w.r.t. parameters
+        # -END- ignore for now -END-
+
+        # data structures for mliap neighbor list
+        # only neighbors strictly inside descriptor cutoff
+
+        int nlocal
+        int nghost
+        int ntotal
+        int * elems
+        int nlistatoms          # current number of atoms in neighborlist
+        int * numneighs         # neighbors count for each atom
+        int * iatoms            # index of each atom
+        int * pair_i            # index of each i atom for each ij pair
+        int * ielems            # element of each atom
+        int nneigh_max          # number of ij neighbors allocated
+        int npairs              # number of ij neighbor pairs
+        int * jatoms            # index of each neighbor
+        int * jelems            # element of each neighbor
+        double ** rij           # distance vector of each neighbor
+        # ----- write only -----
+        double *** graddesc     # descriptor gradient w.r.t. each neighbor
+        # -END- write only -END-
+        int eflag               # indicates if energy is needed
+        int vflag               # indicates if virial is needed
+
+
+cdef extern from "mliap_unified.h" namespace "LAMMPS_NS":
+    cdef cppclass MLIAPDummyDescriptor:
+        MLIAPDummyDescriptor(PyObject *, LAMMPS *)
+        int ndescriptors    # number of descriptors
+        int nelements       # # of unique elements
+        char **elements     # names of unique elements
+        double cutmax       # maximum cutoff needed
+        double rcutfac
+        double *radelem     # element radii
+
+        void compute_descriptors(MLIAPData *)
+        void compute_forces(MLIAPData *)
+        void set_elements(char **, int)
+
+    cdef cppclass MLIAPDummyModel:
+        MLIAPDummyModel(PyObject *, LAMMPS *, char * = NULL)
+        int ndescriptors    # number of descriptors
+        int nparams         # number of parameters per element
+        int nelements;      # # of unique elements
+
+        void compute_gradients(MLIAPData *)
+
+    cdef void update_pair_energy(MLIAPData *, double *)
+    cdef void update_pair_forces(MLIAPData *, double *)
+
+
+# @property sans getter
+def write_only_property(fset):
+    return property(fget=None, fset=fset)
+
+
+cdef class MLIAPDataPy:
+    cdef MLIAPData * data
+
+    def __cinit__(self):
+        self.data = NULL
+ 
+    def update_pair_energy(self, eij):
+        cdef double[:] eij_arr = eij
+        update_pair_energy(self.data, &eij_arr[0])
+    
+    def update_pair_forces(self, fij):
+        cdef double[:, ::1] fij_arr = fij
+        update_pair_forces(self.data, &fij_arr[0][0])
+
+    @property
+    def f(self):
+        if self.data.f is NULL:
+            return None
+        return np.asarray(<double[:self.ntotal, :3]> &self.data.f[0][0])
+    
+    @property
+    def size_gradforce(self):
+        return self.data.size_gradforce
+ 
+    @write_only_property
+    def gradforce(self, value):
+        if self.data.gradforce is NULL:
+            raise ValueError("attempt to set NULL gradforce")
+        cdef double[:, :] gradforce_view = <double[:self.ntotal, :self.size_gradforce]> &self.data.gradforce[0][0]
+        cdef double[:, :] value_view = value
+        gradforce_view[:] = value_view
+ 
+    @write_only_property
+    def betas(self, value):
+        if self.data.betas is NULL:
+            raise ValueError("attempt to set NULL betas")
+        cdef double[:, :] betas_view = <double[:self.nlistatoms, :self.ndescriptors]> &self.data.betas[0][0]
+        cdef double[:, :] value_view = value
+        betas_view[:] = value_view
+
+    @write_only_property
+    def descriptors(self, value):
+        if self.data.descriptors is NULL:
+            raise ValueError("attempt to set NULL descriptors")
+        cdef double[:, :] descriptors_view = <double[:self.nlistatoms, :self.ndescriptors]> &self.data.descriptors[0][0]
+        cdef double[:, :] value_view = value
+        descriptors_view[:] = value_view
+
+    @write_only_property
+    def eatoms(self, value):
+        if self.data.eatoms is NULL:
+            raise ValueError("attempt to set NULL eatoms")
+        cdef double[:] eatoms_view = <double[:self.nlistatoms]> &self.data.eatoms[0]
+        cdef double[:] value_view = value
+        eatoms_view[:] = value_view
+
+    @write_only_property
+    def energy(self, value):
+        self.data.energy = <double>value
+
+    @property
+    def ndescriptors(self):
+        return self.data.ndescriptors
+
+    @property
+    def nparams(self):
+        return self.data.nparams
+
+    @property
+    def nelements(self):
+        return self.data.nelements
+
+    # data structures for grad-grad list (gamma)
+
+    @property
+    def gamma_nnz(self):
+        return self.data.gamma_nnz
+
+    @property
+    def gamma(self):
+        if self.data.gamma is NULL:
+            return None
+        return np.asarray(<double[:self.nlistatoms, :self.gama_nnz]> &self.data.gamma[0][0])
+
+    @property
+    def gamma_row_index(self):
+        if self.data.gamma_row_index is NULL:
+            return None
+        return np.asarray(<int[:self.nlistatoms, :self.gamma_nnz]> &self.data.gamma_row_index[0][0])
+
+    @property
+    def gamma_col_index(self):
+        if self.data.gamma_col_index is NULL:
+            return None
+        return np.asarray(<int[:self.nlistatoms, :self.gamma_nnz]> &self.data.gamma_col_index[0][0])
+
+    @property
+    def egradient(self):
+        if self.data.egradient is NULL:
+            return None
+        return np.asarray(<double[:self.nelements*self.nparams]> &self.data.egradient[0])
+
+    # data structures for mliap neighbor list
+    # only neighbors strictly inside descriptor cutoff
+
+    @property
+    def nlocal(self):
+        return self.data.nlocal
+
+    @property
+    def nghost(self):
+        return self.data.nghost
+
+    @property
+    def ntotal(self):
+        return self.data.ntotal
+    
+    @property
+    def elems(self):
+        if self.data.elems is NULL:
+            return None
+        return np.asarray(<int[:self.ntotal]> &self.data.elems[0])
+
+    @property
+    def nlistatoms(self):
+        return self.data.nlistatoms
+
+    @property
+    def numneighs(self):
+        if self.data.numneighs is NULL:
+            return None
+        return np.asarray(<int[:self.nlistatoms]> &self.data.numneighs[0])
+
+    @property
+    def iatoms(self):
+        if self.data.iatoms is NULL:
+            return None
+        return np.asarray(<int[:self.nlistatoms]> &self.data.iatoms[0])
+    
+    @property
+    def pair_i(self):
+        if self.data.pair_i is NULL:
+            return None
+        return np.asarray(<int[:self.npairs]> &self.data.pair_i[0])
+
+    @property
+    def ielems(self):
+        if self.data.ielems is NULL:
+            return None
+        return np.asarray(<int[:self.nlistatoms]> &self.data.ielems[0])
+    
+    @property
+    def npairs(self):
+        return self.data.npairs
+
+    @property
+    def jatoms(self):
+        if self.data.jatoms is NULL:
+            return None
+        return np.asarray(<int[:self.npairs]> &self.data.jatoms[0])
+    
+    @property
+    def pair_j(self):
+        return self.jatoms
+
+    @property
+    def jelems(self):
+        if self.data.jelems is NULL:
+            return None
+        return np.asarray(<int[:self.npairs]> &self.data.jelems[0])
+
+    @property
+    def rij(self):
+        if self.data.rij is NULL:
+            return None
+        return np.asarray(<double[:self.npairs, :3]> &self.data.rij[0][0])
+
+    @write_only_property
+    def graddesc(self, value):
+        if self.data.graddesc is NULL:
+            raise ValueError("attempt to set NULL graddesc")
+        cdef double[:, :, :] graddesc_view = <double[:self.npairs, :self.ndescriptors, :3]> &self.data.graddesc[0][0][0]
+        cdef double[:, :, :] value_view = value
+        graddesc_view[:] = value_view
+
+    @property
+    def eflag(self):
+        return self.data.eflag
+
+    @property
+    def vflag(self):
+        return self.data.vflag
+
+
+cdef class MLIAPUnifiedInterface:
+    cdef MLIAPDummyModel * model
+    cdef MLIAPDummyDescriptor * descriptor
+    cdef unified_impl
+
+    def __init__(self, unified_impl):
+        self.model = NULL
+        self.descriptor = NULL
+        self.unified_impl = unified_impl
+    
+    def compute_gradients(self, data):
+        self.unified_impl.compute_gradients(data)
+    
+    def compute_descriptors(self, data):
+        self.unified_impl.compute_descriptors(data)
+    
+    def compute_forces(self, data):
+        self.unified_impl.compute_forces(data)
+
+
+cdef public void compute_gradients_python(unified_int, MLIAPData *data) with gil:
+    pydata = MLIAPDataPy()
+    pydata.data = data
+    unified_int.compute_gradients(pydata)
+
+
+cdef public void compute_descriptors_python(unified_int, MLIAPData *data) with gil:
+    pydata = MLIAPDataPy()
+    pydata.data = data
+    unified_int.compute_descriptors(pydata)
+
+
+cdef public void compute_forces_python(unified_int, MLIAPData *data) with gil:
+    pydata = MLIAPDataPy()
+    pydata.data = data
+    unified_int.compute_forces(pydata)
+
+
+cdef public object mliap_unified_connect(char *fname, MLIAPDummyModel * model,
+                                         MLIAPDummyDescriptor * descriptor) with gil:
+    str_fname = fname.decode('utf-8')
+    with open(str_fname, 'rb') as pfile:
+        unified = pickle.load(pfile)
+
+    unified_int = MLIAPUnifiedInterface(unified)
+    unified_int.model = model
+    unified_int.descriptor = descriptor
+
+    unified.interface = unified_int
+
+    if unified.ndescriptors is None:
+        raise ValueError("no descriptors set")
+
+    unified_int.descriptor.ndescriptors = <int>unified.ndescriptors
+    unified_int.descriptor.rcutfac = <double>unified.rcutfac
+    unified_int.model.ndescriptors = <int>unified.ndescriptors
+    unified_int.model.nparams = <int>unified.nparams
+
+    if unified.element_types is None:
+        raise ValueError("no element type set")
+    
+    cdef int nelements = <int>len(unified.element_types)
+    cdef char **elements = <char**>malloc(nelements * sizeof(char*))
+    if not elements:
+        raise MemoryError("failed to allocate memory for element names")
+    cdef char *elem_name
+    for i, elem in enumerate(unified.element_types):
+        elem_name_bytes = elem.encode('UTF-8')
+        elem_name = elem_name_bytes
+        elements[i] = &elem_name[0]
+    unified_int.descriptor.set_elements(elements, nelements)
+    unified_int.model.nelements = nelements
+
+    free(elements)
+    return unified_int

src/ML-IAP/pair_mliap.cpp

@@ -25,6 +25,7 @@
 #include "mliap_model_nn.h"
 #include "mliap_model_quadratic.h"
 #ifdef MLIAP_PYTHON
+#include "mliap_unified.h"
 #include "mliap_model_python.h"
 #endif
 
@@ -93,11 +94,11 @@ void PairMLIAP::compute(int eflag, int vflag)
   // compute E_i and beta_i = dE_i/dB_i for all i in list
 
   model->compute_gradients(data);
-  e_tally(data);
 
   // calculate force contributions beta_i*dB_i/dR_j
 
   descriptor->compute_forces(data);
+  e_tally(data);
 
   // calculate stress
 
@@ -124,7 +125,7 @@ void PairMLIAP::allocate()
 
 void PairMLIAP::settings(int narg, char ** arg)
 {
-  if (narg < 4)
+  if (narg < 2)
     error->all(FLERR,"Illegal pair_style command");
 
   // set flags for required keywords
@@ -174,6 +175,16 @@ void PairMLIAP::settings(int narg, char ** arg)
 
       } else error->all(FLERR,"Illegal pair_style mliap command");
       descriptorflag = 1;
+#ifdef MLIAP_PYTHON
+    } else if (strcmp(arg[iarg], "unified") == 0) {
+      if (iarg+1 > narg) error->all(FLERR,"Illegal pair_style mliap command");
+      MLIAPBuildUnified_t build = build_unified(arg[iarg+1], data, lmp);
+      model = build.model;
+      descriptor = build.descriptor;
+      modelflag = 1;
+      descriptorflag = 1;
+      iarg += 2;
+#endif
     } else
       error->all(FLERR,"Illegal pair_style mliap command");
   }

src/PYTHON/python_impl.cpp

@@ -29,10 +29,12 @@
 
 #ifdef MLIAP_PYTHON
 #include "mliap_model_python.h"
+#include "mliap_unified.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"
+#include "mliap_unifiedpy.h"
 #endif
 
 using namespace LAMMPS_NS;
@@ -66,6 +68,9 @@ PythonImpl::PythonImpl(LAMMPS *lmp) : Pointers(lmp)
   // 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.");
+
+  err = PyImport_AppendInittab("mliap_unifiedpy", PyInit_mliap_unifiedpy);
+  if (err) error->all(FLERR, "Could not register MLIAPPY unified embedded python module.");
 #endif
 
   Py_Initialize();

src/library.cpp

@@ -5576,6 +5576,13 @@ int lammps_get_last_error_message(void *handle, char *buffer, int buf_size) {
   return 0;
 }
 
+#ifdef LMP_PYTHON
+#include <Python.h>
+int lammps_PYTHON_API_VERSION(){
+  return PYTHON_API_VERSION;
+}
+#endif
+
 // Local Variables:
 // fill-column: 72
 // End:

src/library.h

@@ -256,6 +256,10 @@ void lammps_force_timeout(void *handle);
 int lammps_has_error(void *handle);
 int lammps_get_last_error_message(void *handle, char *buffer, int buf_size);
 
+#ifdef LMP_PYTHON
+int lammps_PYTHON_API_VERSION();
+#endif
+
 #ifdef __cplusplus
 }
 #endif