Merge pull request #3458 from Boogie3D/mliappy_unified

MLIAP Unified Interface

python/lammps/mliap/__init__.py

@@ -5,6 +5,7 @@
 import sysconfig
 import ctypes
 import platform
+import warnings
 
 py_ver = sysconfig.get_config_vars('VERSION')[0]
 OS_name = platform.system()
@@ -25,8 +26,10 @@ except Exception as e:
     raise OSError("Unable to locate python shared library") from e
 
 if not pylib.Py_IsInitialized():
-    raise RuntimeError("This interpreter is not compatible with python-based mliap for LAMMPS.")
+    warnings.warn("This interpreter is not compatible with python-based MLIAP for LAMMPS. "
+                  "Attempting to activate the MLIAP-python coupling from python may result "
+                  "in undefined behavior.")
+else:
+    from .loader import load_model, load_unified, activate_mliappy
 
 del sysconfig, ctypes, library, pylib
-
-from .loader import load_model, activate_mliappy

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
-        # 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
+        library = lmp.lib
+        module_names = ["mliap_model_python_couple", "mliap_unified_couple"]
+        api_version = library.lammps_python_api_version()
 
-        path = lmp.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)
-        # End Importlib magic to find the embedded python module
+        for module_name in module_names:
+            # Make Machinery
+            loader = DynamicLoader(module_name,library,api_version)
+            spec = importlib.util.spec_from_loader(module_name,loader)
 
+            # Do the import
+            module = importlib.util.module_from_spec(spec)
+            sys.modules[module_name] = module
+            spec.loader.exec_module(module)
     except Exception as ee:
         raise ImportError("Could not load ML-IAP python coupling module.") from ee
 
@@ -50,3 +81,11 @@ def load_model(model):
                           ) from ie
     mliap_model_python_couple.load_from_python(model)
 
+def load_unified(model):
+    try:
+        import mliap_unified_couple
+    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_unified_couple.load_from_python(model)

python/lammps/mliap/mliap_unified_abc.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+import pickle
+
+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."""
+
+    def pickle(self, fname):
+        with open(fname, 'wb') as fp:
+            pickle.dump(self, fp)

python/lammps/mliap/mliap_unified_lj.py

@@ -0,0 +1,44 @@
+from .mliap_unified_abc import MLIAPUnified
+import numpy as np
+
+
+class MLIAPUnifiedLJ(MLIAPUnified):
+    """Test implementation for MLIAPUnified."""
+
+    def __init__(self, element_types, epsilon=1.0, sigma=1.0, rcutfac=1.25):
+        super().__init__()
+        self.element_types = element_types
+        self.ndescriptors = 1
+        self.nparams = 3
+        # Mimicking the LJ pair-style:
+        # pair_style lj/cut 2.5
+        # pair_coeff * * 1 1
+        self.epsilon = epsilon
+        self.sigma = sigma
+        self.rcutfac = rcutfac
+
+    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

python/lammps/mliap/pytorch.py

@@ -80,10 +80,10 @@ class TorchWrapper(torch.nn.Module):
 
             n_params : torch.nn.Module (None)
                 Number of NN model parameters
-            
+
             device : torch.nn.Module (None)
                 Accelerator device
-            
+
             dtype : torch.dtype (torch.float64)
                 Dtype to use on device
         """
@@ -325,6 +325,6 @@ class ElemwiseModels(torch.nn.Module):
         per_atom_attributes = torch.zeros(elems.size(dim=0), dtype=self.dtype)
         given_elems, elem_indices = torch.unique(elems, return_inverse=True)
         for i, elem in enumerate(given_elems):
-            self.subnets[elem].to(self.dtype) 
+            self.subnets[elem].to(self.dtype)
             per_atom_attributes[elem_indices == i] = self.subnets[elem](descriptors[elem_indices == i]).flatten()
         return per_atom_attributes