Changes to MLIAP python

- update lammps python package to use setuptools
- refactor MLIAP classes into lammps python package

lammps.mliap package
- change TorchWrapper to use dtype and device as arguments
- turn activation of mliappy into functions (was a class)
- add a check to see if python interpreter is compatible
  with python lib calls internal to lammps

mliap_model_python_couple.pyx:
- load models ending in '.pt' or '.pth' with pytorch rather than pickle

doc/src/pair_mliap.rst

@@ -84,6 +84,16 @@ for the :doc:`pair_style snap <pair_snap>` coefficient file.
 Specifically, the line containing the element weight and radius is omitted,
 since these are handled by the *descriptor*.
 
+Notes on mliappy models:
+When the *model* keyword is *mliappy*, the filename should end in '.pt',
+'.pth' for pytorch models, or be a pickle file. To load a model from
+memory (i.e. an existing python object), specify the filename as
+"LATER", and then call `lammps.mliap.load_model(model)` from python
+before using the pair style. When using lammps via the library mode, you will need to call
+`lammps.mliappy.activate_mliappy(lmp)` on the active lammps object
+before the pair style is defined. This call locates and loads the mliap-specific
+python module that is built into lammps.
+
 The *descriptor* keyword is followed by a descriptor style, and additional arguments.
 Currently the only descriptor style is *sna*, indicating the bispectrum component
 descriptors used by the Spectral Neighbor Analysis Potential (SNAP) potentials of

examples/mliap/Ta06A.mliap.pytorch

@@ -11,7 +11,7 @@ variable zblz equal 73
 
 pair_style hybrid/overlay &
 zbl ${zblcutinner} ${zblcutouter} &
-mliap model mliappy Ta06A.mliap.pytorch.model.pkl &
+mliap model mliappy Ta06A.mliap.pytorch.model.pt &
 descriptor sna Ta06A.mliap.descriptor
 pair_coeff 1 1 zbl ${zblz} ${zblz}
 pair_coeff * * mliap Ta

examples/mliap/convert_mliap_Ta06A.py

@@ -1,13 +1,9 @@
 import sys
 import numpy as np
 import torch
-import pickle
-import os
-import shutil
 
-shutil.copyfile('../../src/MLIAP/mliappy_pytorch.py','./mliappy_pytorch.py')
-
-import mliappy_pytorch
+# torch.nn.modules useful for defining a MLIAPPY model.
+from lammps.mliap.pytorch import TorchWrapper, IgnoreElems
 
 # Read coefficients
 coeffs = np.genfromtxt("Ta06A.mliap.model",skip_header=6)
@@ -21,13 +17,10 @@ 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 energy model
-model = mliappy_pytorch.IgnoreElems(lin)
+# Wrap the pytorch model for usage with mliappy coupling.
+model = IgnoreElems(lin) # The linear module does not use the types.
 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)
+n_elements = 1
+linked_model = TorchWrapper(model,n_descriptors=n_descriptors,n_elements=n_elements)
 
-# Save the result.
-with open("Ta06A.mliap.pytorch.model.pkl",'wb') as pfile:
-    pickle.dump(linked_model,pfile)
+torch.save(linked_model,"Ta06A.mliap.pytorch.model.pt")

examples/mliap/mliap_pytorch_Ta06A.py

@@ -81,26 +81,24 @@ import lammps
 
 lmp = lammps.lammps(cmdargs=['-echo','both'])
 
-# this commmand must be run before the MLIAP object is declared in lammps.
+# Before defining the pair style, one must do the following:
+import lammps.mliap
+lammps.mliap.activate_mliappy(lmp)
+# Otherwise, when running lammps in library mode,
+# you will get an error:
+# "ERROR: Loading MLIAPPY coupling module failure."
 
-lmp.mliappy.activate()
-
-# setup the simulation and declare an empty model
+# Setup the simulation and declare an empty model
 # by specifying model filename as "LATER"
-
 lmp.commands_string(before_loading)
 
-# define the PyTorch model by loading a pkl file.
-# this could also be done in other ways.
+# Define the model however you like. In this example
+# we load it from disk:
+import torch
+model = torch.load('Ta06A.mliap.pytorch.model.pt')
 
-import pickle
-with open('Ta06A.mliap.pytorch.model.pkl','rb') as pfile:
-  model = pickle.load(pfile)
-  
-# connect the PyTorch model to the mliap pair style
-
-lmp.mliappy.load_model(model)
+# Connect the PyTorch model to the mliap pair style.
+lammps.mliap.load_model(model)
   
 # run the simulation with the mliap pair style
-
 lmp.commands_string(after_loading)

python/install.py

@@ -95,22 +95,26 @@ print("Installing LAMMPS Python package version %s into site-packages folder" %
 # we need to switch to the folder of the python package
 os.chdir(os.path.dirname(args.package))
 
-from distutils.core import setup
+from setuptools import setup, find_packages
 from distutils.sysconfig import get_python_lib
 import site
-tryuser=False
 
+#Arguments common to global or user install -- everything but data_files
+setup_kwargs= dict(name="lammps",
+        version=verstr,
+        author="Steve Plimpton",
+        author_email="sjplimp@sandia.gov",
+        url="https://lammps.sandia.gov",
+        description="LAMMPS Molecular Dynamics Python package",
+        license="GPL",
+        packages=find_packages(),
+        )
+
+tryuser=False
 try:
   sys.argv = ["setup.py","install"]    # as if had run "python setup.py install"
-  setup(name = "lammps",
-        version = verstr,
-        author = "Steve Plimpton",
-        author_email = "sjplimp@sandia.gov",
-        url = "https://lammps.sandia.gov",
-        description = "LAMMPS Molecular Dynamics Python package",
-        license = "GPL",
-        packages=['lammps'],
-        data_files = [(os.path.join(get_python_lib(), 'lammps'), [args.lib])])
+  setup_kwargs['data_files']=[(os.path.join(get_python_lib(), 'lammps'), [args.lib])]
+  setup(**setup_kwargs)
 except:
   tryuser=True
   print ("Installation into global site-packages folder failed.\nTrying user folder %s now." % site.USER_SITE)
@@ -118,14 +122,7 @@ except:
 if tryuser:
   try:
     sys.argv = ["setup.py","install","--user"]    # as if had run "python setup.py install --user"
-    setup(name = "lammps",
-          version = verstr,
-          author = "Steve Plimpton",
-          author_email = "sjplimp@sandia.gov",
-          url = "https://lammps.sandia.gov",
-          description = "LAMMPS Molecular Dynamics Python package",
-          license = "GPL",
-          packages=['lammps'],
-          data_files = [(os.path.join(site.USER_SITE, 'lammps'), [args.lib])])
+    setup_kwargs['data_files']=[(os.path.join(site.USER_SITE, 'lammps'), [args.lib])]
+    setup(**setup_kwargs)
   except:
     print("Installation into user site package folder failed.")

python/lammps/core.py

@@ -379,8 +379,6 @@ 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
 
@@ -1673,41 +1671,4 @@ class lammps(object):
     idx = self.lib.lammps_find_compute_neighlist(self.lmp, computeid, request)
     return idx
 
-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)

python/lammps/mliap/__init__.py

@@ -0,0 +1,13 @@
+
+# Check compatiblity of this build with the python shared library.
+# If this fails, lammps will segfault because its library will
+# try to improperly start up a new interpreter.
+import sysconfig
+import ctypes
+library = sysconfig.get_config_vars('INSTSONAME')[0]
+pylib = ctypes.CDLL(library)
+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

python/lammps/mliap/loader.py

@@ -0,0 +1,52 @@
+# ----------------------------------------------------------------------
+#   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.
+# -------------------------------------------------------------------------
+
+# ----------------------------------------------------------------------
+#   Contributing author: Nicholas Lubbers (LANL)
+# -------------------------------------------------------------------------
+
+
+import sys
+import importlib.util
+import importlib.machinery
+
+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
+
+        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
+
+    except Exception as ee:
+        raise ImportError("Could not load MLIAP python coupling module.") from ee
+
+def load_model(model):
+    try:
+        import mliap_model_python_couple
+    except ImportError as ie:
+        raise ImportError("MLIAP python module must be activated before loading\n"
+                          "the pair style. Call lammps.mliap.activate_mliappy(lmp)."
+                          ) from ie
+    mliap_model_python_couple.load_from_python(model)
+

python/lammps/mliap/pytorch.py

@@ -22,20 +22,28 @@ 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):
+    def __init__(self, model,n_descriptors,n_elements,n_params=None,device=None,dtype=torch.float64):
         super().__init__()
+
         self.model = model
-        self.model.to(self.dtype)
+        self.device = device
+        self.dtype = dtype
+
+        # Put model on device and convert to dtype
+        self.to(self.dtype)
+        self.to(self.device)
+
         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):
+    def forward(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
+        bispectrum = torch.from_numpy(bispectrum).to(dtype=self.dtype, device=self.device).requires_grad_(True)
+        elems = torch.from_numpy(elems).to(dtype=torch.long, device=self.device) - 1
 
         with torch.autograd.enable_grad():
 
@@ -48,12 +56,6 @@ class TorchWrapper(torch.nn.Module):
         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__()

python/setup.py

@@ -1,6 +1,6 @@
 # this only installs the LAMMPS python package
 # it assumes the LAMMPS shared library is already installed
-from distutils.core import setup
+from setuptools import setup, find_packages
 import os
 
 LAMMPS_PYTHON_DIR = os.path.dirname(os.path.realpath(__file__))
@@ -22,5 +22,5 @@ setup(
     url = "https://lammps.sandia.gov",
     description = "LAMMPS Molecular Dynamics Python package",
     license = "GPL",
-    packages=["lammps"]
+    packages=find_packages(),
 )

src/MLIAP/mliap_model_python_couple.pyx

@@ -46,7 +46,7 @@ cdef object c_id(MLIAPModelPython * c_model):
     """
     return int(<uintptr_t> c_model)
 
-cdef object retrieve(MLIAPModelPython * c_model):
+cdef object retrieve(MLIAPModelPython * c_model) with gil:
     try:
         model = LOADED_MODELS[c_id(c_model)]
     except KeyError as ke:
@@ -61,8 +61,12 @@ cdef public int MLIAPPY_load_model(MLIAPModelPython * c_model, char* fname) with
         model = None
         returnval = 0
     else:
-        with open(str_fname,'rb') as pfile:
-            model = pickle.load(pfile)
+        if str_fname.endswith(".pt") or str_fname.endswith('.pth'):
+            import torch
+            model = torch.load(str_fname)
+        else:
+            with open(str_fname,'rb') as pfile:
+                model = pickle.load(pfile)
         returnval = 1
     LOADED_MODELS[c_id(c_model)] = model
     return returnval