First draft of MLPotential

LICENSE

@@ -1,6 +1,8 @@
 MIT License
 
-Copyright (c) 2021 OpenMM
+Copyright (c) 2021 Stanford University and the Authors.
+Authors: Peter Eastman
+Contributors:
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

openmmml/__init__.py

@@ -0,0 +1,2 @@
+from .mlpotential import MLPotential
+from . import models

openmmml/mlpotential.py

@@ -0,0 +1,67 @@
+"""
+mlpotential.py: Provides a common API for creating OpenMM Systems with ML potentials.
+
+This is part of the OpenMM molecular simulation toolkit originating from
+Simbios, the NIH National Center for Physics-Based Simulation of
+Biological Structures at Stanford, funded under the NIH Roadmap for
+Medical Research, grant U54 GM072970. See https://simtk.org.
+
+Portions copyright (c) 2021 Stanford University and the Authors.
+Authors: Peter Eastman
+Contributors:
+
+Permission is hereby granted, free of charge, to any person obtaining a
+copy of this software and associated documentation files (the "Software"),
+to deal in the Software without restriction, including without limitation
+the rights to use, copy, modify, merge, publish, distribute, sublicense,
+and/or sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
+USE OR OTHER DEALINGS IN THE SOFTWARE.
+"""
+
+import openmm
+import openmm.app
+from typing import Dict
+
+
+class MLPotentialImplFactory(object):
+
+    def createImpl(self, name: str, **args) -> "MLPotentialImpl":
+        raise NotImplementedError('Subclasses must implement createImpl()')
+
+
+class MLPotentialImpl(object):
+
+    def addForces(self, topology: openmm.app.Topology, system: openmm.System, **args):
+        raise NotImplementedError('Subclasses must implement addForces()')
+
+
+class MLPotential(object):
+
+    _implFactories: Dict[str, MLPotentialImplFactory] = {}
+
+    def __init__(self, name: str, **args):
+        self._impl = _implFactories[name].createImpl(name, **args)
+
+    def createSystem(self, topology: openmm.app.Topology, **args) -> openmm.System:
+        system = openmm.System()
+        for atom in topology.atoms():
+            if atom.element is None:
+                system.addParticle(0)
+            else:
+                system.addParticle(atom.element.mass)
+        return self._impl.addForces(topology, system, **args)
+
+    @staticmethod
+    def registerImplFactory(name: str, factory: MLPotentialImplFactory):
+        MLPotential._implFactories[name] = factory

openmmml/models/__init__.py

@@ -0,0 +1 @@
+from . import anipotential

openmmml/models/anipotential.py

@@ -0,0 +1,50 @@
+"""
+anipotential.py: Implements the ANI potential function using TorchANI.
+
+This is part of the OpenMM molecular simulation toolkit originating from
+Simbios, the NIH National Center for Physics-Based Simulation of
+Biological Structures at Stanford, funded under the NIH Roadmap for
+Medical Research, grant U54 GM072970. See https://simtk.org.
+
+Portions copyright (c) 2021 Stanford University and the Authors.
+Authors: Peter Eastman
+Contributors:
+
+Permission is hereby granted, free of charge, to any person obtaining a
+copy of this software and associated documentation files (the "Software"),
+to deal in the Software without restriction, including without limitation
+the rights to use, copy, modify, merge, publish, distribute, sublicense,
+and/or sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
+USE OR OTHER DEALINGS IN THE SOFTWARE.
+"""
+
+from openmmml.mlpotential import MLPotential, MLPotentialImpl, MLPotentialImplFactory
+import openmm
+
+class ANIPotentialImplFactory(MLPotentialImplFactory):
+
+    def createImpl(self, name: str, **args) -> MLPotentialImpl:
+        return ANIPotentialImpl(name)
+
+
+class ANIPotentialImpl(MLPotentialImpl):
+
+    def __init__(self, name):
+        self.name = name
+
+    def addForces(self, topology: openmm.app.Topology, system: openmm.System, **args):
+        pass
+
+MLPotential.registerImplFactory('ani1ccx', ANIPotentialImplFactory())
+MLPotential.registerImplFactory('ani2x', ANIPotentialImplFactory())

setup.py

@@ -0,0 +1,42 @@
+"""OpenMM-ML: A high level API for using machine learning models in OpenMM simulations
+"""
+
+from setuptools import setup, find_packages
+
+DOCLINES = __doc__.split("\n")
+
+########################
+__version__ = '1.0'
+VERSION = __version__
+ISRELEASED = False
+########################
+CLASSIFIERS = """\
+Development Status :: 2 - Pre-Alpha
+Intended Audience :: Science/Research
+Intended Audience :: Developers
+License :: OSI Approved :: MIT License
+Programming Language :: Python
+Programming Language :: Python :: 3
+Topic :: Scientific/Engineering :: Bio-Informatics
+Topic :: Scientific/Engineering :: Chemistry
+Operating System :: Microsoft :: Windows
+Operating System :: POSIX
+Operating System :: Unix
+Operating System :: MacOS
+"""
+
+
+setup(
+    name='openmmml',
+    author='Peter Eastman',
+    description=DOCLINES[0],
+    long_description="\n".join(DOCLINES[2:]),
+    version=__version__,
+    license='MIT',
+    url='https://github.com/openmm/openmm-ml',
+    platforms=['Linux', 'Mac OS-X', 'Unix', 'Windows'],
+    classifiers=CLASSIFIERS.splitlines(),
+    packages=find_packages(),
+    zip_safe=False,
+    install_requires=['numpy', 'openmm >= 7.5'])
+