From e7fa0a6bac75824361c9b13d9aecc8667a41fc43 Mon Sep 17 00:00:00 2001 From: Nicholas Lubbers Date: Mon, 21 Dec 2020 11:51:10 -0700 Subject: [PATCH] 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 | 10 ++++ examples/mliap/Ta06A.mliap.pytorch | 2 +- examples/mliap/convert_mliap_Ta06A.py | 21 +++----- examples/mliap/mliap_pytorch_Ta06A.py | 28 +++++----- python/install.py | 37 ++++++------- python/lammps/core.py | 39 -------------- python/lammps/mliap/__init__.py | 13 +++++ python/lammps/mliap/loader.py | 52 +++++++++++++++++++ .../lammps/mliap/pytorch.py | 24 +++++---- python/setup.py | 4 +- src/MLIAP/mliap_model_python_couple.pyx | 10 ++-- 11 files changed, 135 insertions(+), 105 deletions(-) create mode 100644 python/lammps/mliap/__init__.py create mode 100644 python/lammps/mliap/loader.py rename src/MLIAP/mliappy_pytorch.py => python/lammps/mliap/pytorch.py (80%) diff --git a/doc/src/pair_mliap.rst b/doc/src/pair_mliap.rst index 9dd8c5c6b1..af0cc9e855 100644 --- a/doc/src/pair_mliap.rst +++ b/doc/src/pair_mliap.rst @@ -84,6 +84,16 @@ for the :doc:`pair_style 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 diff --git a/examples/mliap/Ta06A.mliap.pytorch b/examples/mliap/Ta06A.mliap.pytorch index c6c6bc653d..5489688a82 100644 --- a/examples/mliap/Ta06A.mliap.pytorch +++ b/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 diff --git a/examples/mliap/convert_mliap_Ta06A.py b/examples/mliap/convert_mliap_Ta06A.py index ba95678c9e..8a7466743f 100644 --- a/examples/mliap/convert_mliap_Ta06A.py +++ b/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") diff --git a/examples/mliap/mliap_pytorch_Ta06A.py b/examples/mliap/mliap_pytorch_Ta06A.py index 71c89f112e..c4387e21a3 100644 --- a/examples/mliap/mliap_pytorch_Ta06A.py +++ b/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. +lammps.mliap.load_model(model) -# connect the PyTorch model to the mliap pair style - -lmp.mliappy.load_model(model) - # run the simulation with the mliap pair style - lmp.commands_string(after_loading) diff --git a/python/install.py b/python/install.py index a6b69c1ee6..8c632c096e 100644 --- a/python/install.py +++ b/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.") diff --git a/python/lammps/core.py b/python/lammps/core.py index e30f036e1b..a75a02e358 100644 --- a/python/lammps/core.py +++ b/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) diff --git a/python/lammps/mliap/__init__.py b/python/lammps/mliap/__init__.py new file mode 100644 index 0000000000..06c127055e --- /dev/null +++ b/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 \ No newline at end of file diff --git a/python/lammps/mliap/loader.py b/python/lammps/mliap/loader.py new file mode 100644 index 0000000000..b875d8d834 --- /dev/null +++ b/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) + diff --git a/src/MLIAP/mliappy_pytorch.py b/python/lammps/mliap/pytorch.py similarity index 80% rename from src/MLIAP/mliappy_pytorch.py rename to python/lammps/mliap/pytorch.py index ff80720ffa..aa2cf1c97c 100644 --- a/src/MLIAP/mliappy_pytorch.py +++ b/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__() diff --git a/python/setup.py b/python/setup.py index 9be04138d5..50e215cb09 100644 --- a/python/setup.py +++ b/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(), ) diff --git a/src/MLIAP/mliap_model_python_couple.pyx b/src/MLIAP/mliap_model_python_couple.pyx index 49bdd47412..1f04964ef6 100644 --- a/src/MLIAP/mliap_model_python_couple.pyx +++ b/src/MLIAP/mliap_model_python_couple.pyx @@ -46,7 +46,7 @@ cdef object c_id(MLIAPModelPython * c_model): """ return int( 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