Implement python-side unified support, add example

examples/mliap/mliap_unified_lj_Ar.py

@@ -0,0 +1,65 @@
+# Demonstrate how to load a unified model from python.
+# This is essentially the same as in.mliap.unified.lj.Ar
+# except that python is the driving program, and lammps
+# is in library mode.
+
+before_loading =\
+"""# 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
+"""
+after_loading =\
+"""
+
+pair_style	mliap unified EXISTS
+pair_coeff	* * Ar
+
+neighbor	0.3 bin
+neigh_modify	every 20 delay 0 check no
+
+fix		1 all nve
+
+thermo		50
+run		250
+"""
+
+import lammps
+
+lmp = lammps.lammps(cmdargs=['-echo','both'])
+
+# Before defining the pair style, you must do the following:
+import lammps.mliap
+lammps.mliap.activate_mliappy(lmp)
+# Otherwise, when running lammps in library mode,
+# you will get an error
+
+# Setup the simulation to just before declaring
+# the mliap unified pair style
+lmp.commands_string(before_loading)
+
+# Define the model however you like. In this example
+# we simply import the unified L-J example from mliap
+from lammps.mliap.mliap_unified_lj import MLIAPUnifiedLJ
+unified = MLIAPUnifiedLJ()
+
+# You can also load the model from a pickle file.
+# import pickle
+# with open('mliap_unified_lj_Ar.pkl', 'rb') as pfile:
+#     unified = pickle.load(pfile)
+
+# Connect the L-J model to the mliap unified pair style.
+lammps.mliap.load_unified(unified)
+
+
+# Run the simulation with the mliap unified pair style
+# Use pre-loaded model by specifying model filename as "EXISTS"
+lmp.commands_string(after_loading)

python/lammps/mliap/__init__.py

@@ -17,4 +17,4 @@ if not pylib.Py_IsInitialized():
     raise RuntimeError("This interpreter is not compatible with python-based mliap for LAMMPS.")
 del sysconfig, ctypes, library, pylib
 
-from .loader import load_model, activate_mliappy
+from .loader import load_model, load_unified, activate_mliappy

python/lammps/mliap/loader.py

@@ -80,3 +80,12 @@ def load_model(model):
                           "the pair style. Call lammps.mliap.activate_mliappy(lmp)."
                           ) from ie
     mliap_model_python_couple.load_from_python(model)
+
+def load_unified(model):
+    try:
+        import mliap_unifiedpy
+    except ImportError as ie:
+        raise ImportError("ML-IAP python module must be activated before loading\n"
+                          "the pair style. Call lammps.mliap.activate_mliappy(lmp)."
+                          ) from ie
+    mliap_unifiedpy.load_from_python(model)

src/ML-IAP/mliap_unifiedpy.pyx

@@ -98,6 +98,9 @@ cdef extern from "mliap_unified.h" namespace "LAMMPS_NS":
     cdef void update_pair_forces(MLIAPData *, double *) except +
 
 
+LOADED_MODEL = None
+
+
 # @property sans getter
 def write_only_property(fset):
     return property(fget=None, fset=fset)
@@ -334,6 +337,14 @@ cdef public void compute_forces_python(unified_int, MLIAPData *data) except * wi
 cdef public object mliap_unified_connect(char *fname, MLIAPDummyModel * model,
                                          MLIAPDummyDescriptor * descriptor) with gil:
     str_fname = fname.decode('utf-8')
+    if str_fname == 'EXISTS':
+        if LOADED_MODEL is None:
+            raise ValueError("No unified model loaded")
+        unified = LOADED_MODEL
+    elif str_fname.endswith(".pt") or str_fname.endswith('.pth'):
+        import torch
+        unified = torch.load(str_fname)
+    else:
         with open(str_fname, 'rb') as pfile:
             unified = pickle.load(pfile)
 
@@ -370,3 +381,8 @@ cdef public object mliap_unified_connect(char *fname, MLIAPDummyModel * model,
 
     free(elements)
     return unified_int
+
+
+def load_from_python(unified):
+    global LOADED_MODEL
+    LOADED_MODEL = unified