From 7a76a6ee367243a2129f54791bec929b84147241 Mon Sep 17 00:00:00 2001
From: Steven Ray Anaya <sanaya@lanl.gov>
Date: Tue, 16 Aug 2022 02:35:31 -0600
Subject: [PATCH] Implement python-side unified support, add example

---
 examples/mliap/mliap_unified_lj_Ar.py | 65 +++++++++++++++++++++++++++
 python/lammps/mliap/__init__.py       |  2 +-
 python/lammps/mliap/loader.py         |  9 ++++
 src/ML-IAP/mliap_unifiedpy.pyx        | 20 ++++++++-
 4 files changed, 93 insertions(+), 3 deletions(-)
 create mode 100644 examples/mliap/mliap_unified_lj_Ar.py

diff --git a/examples/mliap/mliap_unified_lj_Ar.py b/examples/mliap/mliap_unified_lj_Ar.py
new file mode 100644
index 0000000000..9b1c11be7d
--- /dev/null
+++ b/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)
diff --git a/python/lammps/mliap/__init__.py b/python/lammps/mliap/__init__.py
index 57fe97d803..b29346f255 100644
--- a/python/lammps/mliap/__init__.py
+++ b/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
diff --git a/python/lammps/mliap/loader.py b/python/lammps/mliap/loader.py
index e50ba8cb93..284af01b05 100644
--- a/python/lammps/mliap/loader.py
+++ b/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)
diff --git a/src/ML-IAP/mliap_unifiedpy.pyx b/src/ML-IAP/mliap_unifiedpy.pyx
index d34137e3f6..4a488e3eb2 100644
--- a/src/ML-IAP/mliap_unifiedpy.pyx
+++ b/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,8 +337,16 @@ 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')
-    with open(str_fname, 'rb') as pfile:
-        unified = pickle.load(pfile)
+    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)
 
     unified_int = MLIAPUnifiedInterface(unified)
     unified_int.model = model
@@ -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