openmmml 1.6__tar.gz

openmmml-1.6/LICENSE

@@ -0,0 +1,23 @@
+MIT License
+
+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 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.

openmmml-1.6/PKG-INFO

@@ -0,0 +1,34 @@
+Metadata-Version: 2.4
+Name: openmmml
+Version: 1.6
+Summary: OpenMM-ML: A high level API for using machine learning models in OpenMM simulations
+Home-page: https://github.com/openmm/openmm-ml
+Author: Peter Eastman
+License: MIT
+Platform: Linux
+Platform: Mac OS-X
+Platform: Unix
+Platform: Windows
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX
+Classifier: Operating System :: Unix
+Classifier: Operating System :: MacOS
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: openmm>=8.5
+Dynamic: author
+Dynamic: classifier
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: platform
+Dynamic: requires-dist
+Dynamic: summary

openmmml-1.6/README.md

@@ -0,0 +1,83 @@
+# OpenMM-ML
+
+This is a high level API for using machine learning models in OpenMM simulations.  With just a few lines of code, you
+can set up a simulation that uses a standard, pretrained model to represent some or all of the interactions in a system.
+
+### Supported models and frameworks
+
+The current release of OpenMM-ML supports the following potential functions.
+
+- [MACE](https://arxiv.org/abs/2206.07697) models, including the pre-trained [MACE-OFF23](https://arxiv.org/abs/2312.15211),
+  [MACE-MPA-0](https://github.com/ACEsuit/mace-foundations), [MACE-OMAT-0](https://github.com/ACEsuit/mace-foundations),
+  and MACE-OMOL-0 models, utilizing the [MACE implementation](https://github.com/ACEsuit/mace).
+
+- The pretrained [AIMNet2](https://doi.org/10.1039/D4SC08572H) model.
+
+- Models generated with [TorchMD-Net](https://github.com/torchmd/torchmd-net), including the pretrained [AceFF](https://arxiv.org/abs/2601.00581) models.
+
+- Models generated with [FeNNol](https://github.com/FeNNol-tools/FeNNol), including the pretrained [FeNNix-Bio1](https://chemrxiv.org/doi/full/10.26434/chemrxiv-2025-f1hgn-v4) models.
+
+- Pretrained [Orb](https://github.com/orbital-materials/orb-models) models.
+
+- Models generated by [NequIP](https://github.com/mir-group/nequip). Models generated by [Allegro](https://github.com/mir-group/allegro) are also supported by this interface.
+
+- The pretrained [ANI-1ccx](https://www.nature.com/articles/s41467-019-10827-4) and [ANI-2x](https://pubs.acs.org/doi/full/10.1021/acs.jctc.0c00121)
+  models, using the implementations in [TorchANI](https://github.com/aiqm/torchani).
+
+- Models generated by [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit).
+
+In addition to its built in support for the models listed above, OpenMM-ML can use any model implemented with an
+[ASE Calculator](https://ase-lib.org/ase/calculators/calculators.html).  A very wide range of machine learning and quantum
+chemistry packages provide interfaces in this format, allowing them to be used with OpenMM-ML.
+
+See the documentation for details on using each of these model types.
+
+### Documentation 
+
+You can find the current documentation at https://openmm.github.io/openmm-ml.
+
+### Installation
+
+The stable version of OpenMM-ML can be installed with pip.
+
+```
+pip install openmmml
+```
+
+This installs only OpenMM-ML itself, not the packages that provide specific models.  Those must be installed separately
+by following the instructions from the package developers.
+
+OpenMM-ML can also be installed with conda or mamba.
+
+```bash
+mamba install -c conda-forge openmm-ml
+```
+
+However, we generally recommend using pip because most of the packages that provide models are only available
+with pip, not conda.
+
+### Usage
+
+To use this package, create a MLPotential object, specifying the name of the potential function to use.  You can then
+call createSystem() to create a System object  for a simulation.  For example,
+
+```python
+from openmmml import MLPotential
+potential = MLPotential('ani2x')
+system = potential.createSystem(topology)
+```
+
+Alternatively, you can use createMixedSystem() to create a System where part is modeled with this potential and the rest
+is modeled with a conventional force field.  As an example, suppose the Topology contains three chains.  Chain 0 is a
+protein, chain 1 is a ligand, and chain 2 is solvent.  The following code creates a System in which the internal energy
+of the ligand is computed with ANI2x, while everything else (including interactions between the ligand and the rest of
+the System) is computed with Amber14.
+
+```python
+forcefield = ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')
+mm_system = forcefield.createSystem(topology)
+chains = list(topology.chains())
+ml_atoms = [atom.index for atom in chains[1].atoms()]
+potential = MLPotential('ani2x')
+ml_system = potential.createMixedSystem(topology, mm_system, ml_atoms)
+```

openmmml-1.6/openmmml/__init__.py

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

openmmml-1.6/openmmml/mlpotential.py

@@ -0,0 +1,479 @@
+"""
+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-2026 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
+import openmm.unit as unit
+from copy import deepcopy
+from typing import Dict, Iterable, Optional
+import os
+import shutil
+import tempfile
+import urllib.request
+import sys
+if sys.version_info < (3, 10):
+    from importlib_metadata import entry_points
+else:
+    from importlib.metadata import entry_points
+
+
+class MLPotentialImplFactory(object):
+    """Abstract interface for classes that create MLPotentialImpl objects.
+
+    If you are defining a new potential function, you need to create subclasses
+    of MLPotentialImpl and MLPotentialImplFactory, and register an instance of
+    the factory by calling MLPotential.registerImplFactory().  Alternatively,
+    if a Python package creates an entry point in the group "openmmml.potentials",
+    the potential will be registered automatically.  The entry point name is the
+    name of the potential function, and the value should be the name of the
+    MLPotentialImplFactory subclass.
+    """
+    
+    def createImpl(self, name: str, **args) -> "MLPotentialImpl":
+        """Create a MLPotentialImpl that will be used to implement a MLPotential.
+
+        When a MLPotential is created, it invokes this method to create an object
+        implementing the requested potential.  Subclasses must implement this method
+        to return an instance of the correct MLPotentialImpl subclass.
+
+        Parameters
+        ----------
+        name: str
+            the name of the potential that was specified to the MLPotential constructor
+        args:
+            any additional keyword arguments that were provided to the MLPotential
+            constructor are passed to this method.  This allows subclasses to customize
+            their behavior based on extra arguments.
+
+        Returns
+        -------
+        a MLPotentialImpl that implements the potential
+        """
+        raise NotImplementedError('Subclasses must implement createImpl()')
+
+
+class MLPotentialImpl(object):
+    """Abstract interface for classes that implement potential functions.
+
+    If you are defining a new potential function, you need to create subclasses
+    of MLPotentialImpl and MLPotentialImplFactory.  When a user creates a
+    MLPotential and specifies a name for the potential to use, it looks up the
+    factory that has been registered for that name and uses it to create a
+    MLPotentialImpl of the appropriate subclass.
+    """
+    
+    def addForces(self,
+                  topology: openmm.app.Topology,
+                  system: openmm.System,
+                  atoms: Optional[Iterable[int]],
+                  forceGroup: int,
+                  **args):
+        """Add Force objects to a System to implement the potential function.
+
+        This is invoked by MLPotential.createSystem().  Subclasses must implement
+        it to create the requested potential function.
+
+        Parameters
+        ----------
+        topology: Topology
+            the Topology from which the System is being created
+        system: System
+            the System that is being created
+        atoms: Optional[Iterable[int]]
+            the indices of atoms the potential should be applied to, or None if
+            it should be applied to the entire System
+        forceGroup: int
+            the force group that any newly added Forces should be in
+        args:
+            any additional keyword arguments that were provided to createSystem()
+            are passed to this method.  This allows subclasses to customize their
+            behavior based on extra arguments.
+        """
+        raise NotImplementedError('Subclasses must implement addForces()')
+
+    def _getTorchDevice(self, args):
+        """This is a utility routine for use by subclasses that are implemented with PyTorch.  It selects what device
+        to use, which can either by specified by the user with the 'device' argument, or chosen automatically based on
+        the available hardware."""
+        import torch
+        if 'device' in args:
+            device = args['device']
+            if isinstance(device, str):
+                device = torch.device(device)
+            return device
+        if torch.cuda.is_available():
+            return torch.device('cuda')
+        return torch.device('cpu')
+
+    def _getCacheDir(self):
+        """This is a utility routine returning a cache directory for use by subclasses that need to download their own
+        models.  Before returning, it will create the directory and any parent directories if they do not already exist."""
+        path = os.path.expanduser("~/.cache/openmm-ml")
+        os.makedirs(path, exist_ok=True)
+        return path
+
+    def _downloadOrFindFile(self, name: str, url: str):
+        """Downloads a file at the requested URL and saves it in a cache directory under the specified name.  If a file
+        with the specified name already exists in the cache, the path to it will be returned without downloading it again."""
+        cacheDir = self._getCacheDir()
+        targetPath = os.path.join(cacheDir, name)
+        if not os.path.isfile(targetPath):
+            with urllib.request.urlopen(url) as remoteFile:
+                # Download into a temporary directory first so that we do not
+                # end up with a partially downloaded file at targetPath and we
+                # do not create temporary files directly in the cache directory.
+                with tempfile.TemporaryDirectory(dir=cacheDir) as tempDir:
+                    tempPath = os.path.join(tempDir, name)
+                    with open(tempPath, "wb") as localFile:
+                        shutil.copyfileobj(remoteFile, localFile)
+                    os.replace(tempPath, targetPath)
+        return targetPath
+
+class MLPotential(object):
+    """A potential function that can be used in simulations.
+
+    To use this class, create a MLPotential, specifying the name of the potential
+    function to use.  You can then call createSystem() to create a System object
+    for a simulation.  For example,
+
+    >>> potential = MLPotential('ani2x')
+    >>> system = potential.createSystem(topology)
+
+    Alternatively, you can use createMixedSystem() to create a System where part is
+    modeled with this potential and the rest is modeled with a conventional force
+    field.  As an example, suppose the Topology contains three chains.  Chain 0 is
+    a protein, chain 1 is a ligand, and chain 2 is solvent.  The following code
+    creates a System in which the internal energy of the ligand is computed with
+    ANI2x, while everything else (including interactions between the ligand and the
+    rest of the System) is computed with Amber14.
+
+    >>> forcefield = ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')
+    >>> mm_system = forcefield.createSystem(topology)
+    >>> chains = list(topology.chains())
+    >>> ml_atoms = [atom.index for atom in chains[1].atoms()]
+    >>> potential = MLPotential('ani2x')
+    >>> ml_system = potential.createMixedSystem(topology, mm_system, ml_atoms)
+    """
+
+    _implFactories: Dict[str, MLPotentialImplFactory] = {}
+    
+    def __init__(self, name: str, **args):
+        """Create a MLPotential.
+
+        Parameters
+        ----------
+        name: str
+            the name of the potential function to use.  Built in support is currently
+            provided for the following: 'ani1ccx', 'ani2x'.  Others may be added by
+            calling MLPotential.registerImplFactory().
+        args:
+            particular potential functions may define additional arguments that can
+            be used to customize them.  See the documentation on the specific
+            potential functions for more information.
+        """
+        self._impl = MLPotential._implFactories[name].createImpl(name, **args)
+
+    def createSystem(self, topology: openmm.app.Topology, removeCMMotion: bool = True, **args) -> openmm.System:
+        """Create a System for running a simulation with this potential function.
+
+        Parameters
+        ----------
+        topology: Topology
+            the Topology for which to create a System
+        removeCMMotion: bool
+            if true, a CMMotionRemover will be added to the System. 
+        args:
+            particular potential functions may define additional arguments that can
+            be used to customize them.  See the documentation on the specific
+            potential functions for more information.
+
+        Returns
+        -------
+        a newly created System object that uses this potential function to model the Topology
+        """
+        system = openmm.System()
+        if topology.getPeriodicBoxVectors() is not None:
+            system.setDefaultPeriodicBoxVectors(*topology.getPeriodicBoxVectors())
+        for atom in topology.atoms():
+            if atom.element is None:
+                system.addParticle(0)
+            else:
+                system.addParticle(atom.element.mass)
+        self._impl.addForces(topology, system, None, 0, **args)
+        if removeCMMotion:
+            system.addForce(openmm.CMMotionRemover())
+        return system
+
+    def createMixedSystem(self,
+                          topology: openmm.app.Topology,
+                          system: openmm.System,
+                          atoms: Iterable[int],
+                          removeConstraints: bool = True,
+                          forceGroup: int = 0,
+                          interpolate: bool = False,
+                          **args) -> openmm.System:
+        """Create a System that is partly modeled with this potential and partly
+        with a conventional force field.
+
+        To use this method, first create a System that is entirely modeled with the
+        conventional force field.  Pass it to this method, along with the indices of the
+        atoms to model with this potential (the "ML subset").  It returns a new System
+        that is identical to the original one except for the following changes.
+
+        1. Removing all bonds, angles, and torsions for which *all* atoms are in the
+           ML subset.
+        2. For every NonbondedForce and CustomNonbondedForce, adding exceptions/exclusions
+           to prevent atoms in the ML subset from interacting with each other.
+        3. (Optional) Removing constraints between atoms that are both in the ML subset.
+        4. Adding Forces as necessary to compute the internal energy of the ML subset
+           with this potential.
+
+        Alternatively, the System can include Forces to compute the energy both with the
+        conventional force field and with this potential, and to smoothly interpolate
+        between them.  In that case, it creates a CustomCVForce containing the following.
+
+        1. The Forces to compute this potential.
+        2. Forces to compute the bonds, angles, and torsions that were removed above.
+        3. For every NonbondedForce, a corresponding CustomBondForce to compute the
+           nonbonded interactions within the ML subset.
+
+        The CustomCVForce defines a global parameter called "lambda_interpolate" that interpolates
+        between the two potentials.  When lambda_interpolate=0, the energy is computed entirely with
+        the conventional force field.  When lambda_interpolate=1, the energy is computed entirely with
+        the ML potential.  You can set its value by calling setParameter() on the Context.
+
+        Parameters
+        ----------
+        topology: Topology
+            the Topology for which to create a System
+        system: System
+            a System that models the Topology with a conventional force field
+        atoms: Iterable[int]
+            the indices of all atoms whose interactions should be computed with
+            this potential
+        removeConstraints: bool
+            if True, remove constraints between pairs of atoms whose interaction
+            will be computed with this potential
+        forceGroup: int
+            the force group the ML potential's Forces should be placed in
+        interpolate: bool
+            if True, create a System that can smoothly interpolate between the conventional
+            and ML potentials
+        args:
+            particular potential functions may define additional arguments that can
+            be used to customize them.  See the documentation on the specific
+            potential functions for more information.
+
+        Returns
+        -------
+        a newly created System object that uses this potential function to model the Topology
+        """
+        # Create the new System, removing bonded interactions within the ML subset.
+
+        newSystem = self._removeBonds(system, atoms, True, removeConstraints)
+
+        # Add nonbonded exceptions and exclusions.
+
+        atomList = list(atoms)
+        for force in newSystem.getForces():
+            if isinstance(force, openmm.NonbondedForce):
+                for i in range(len(atomList)):
+                    for j in range(i):
+                        force.addException(atomList[i], atomList[j], 0, 1, 0, True)
+            elif isinstance(force, openmm.CustomNonbondedForce):
+                existing = set(tuple(force.getExclusionParticles(i)) for i in range(force.getNumExclusions()))
+                for i in range(len(atomList)):
+                    a1 = atomList[i]
+                    for j in range(i):
+                        a2 = atomList[j]
+                        if (a1, a2) not in existing and (a2, a1) not in existing:
+                            force.addExclusion(a1, a2)
+
+        # Add the ML potential.
+
+        if not interpolate:
+            self._impl.addForces(topology, newSystem, atomList, forceGroup, **args)
+        else:
+            # Create a CustomCVForce and put the ML forces inside it.
+
+            cv = openmm.CustomCVForce('')
+            cv.addGlobalParameter('lambda_interpolate', 1)
+            tempSystem = openmm.System()
+            self._impl.addForces(topology, tempSystem, atomList, forceGroup, **args)
+            mlVarNames = []
+            for i, force in enumerate(tempSystem.getForces()):
+                name = f'mlForce{i+1}'
+                cv.addCollectiveVariable(name, deepcopy(force))
+                mlVarNames.append(name)
+
+            # Create Forces for all the bonded interactions within the ML subset and add them to the CustomCVForce.
+
+            bondedSystem = self._removeBonds(system, atoms, False, removeConstraints)
+            bondedForces = []
+            for force in bondedSystem.getForces():
+                if hasattr(force, 'addBond') or hasattr(force, 'addAngle') or hasattr(force, 'addTorsion'):
+                    bondedForces.append(force)
+            mmVarNames = []
+            for i, force in enumerate(bondedForces):
+                name = f'mmForce{i+1}'
+                cv.addCollectiveVariable(name, deepcopy(force))
+                mmVarNames.append(name)
+
+            # Create a CustomBondForce that computes all nonbonded interactions within the ML subset.
+
+            for force in system.getForces():
+                if isinstance(force, openmm.NonbondedForce):
+                    internalNonbonded = openmm.CustomBondForce('138.935456*chargeProd/r + 4*epsilon*((sigma/r)^12-(sigma/r)^6)')
+                    internalNonbonded.addPerBondParameter('chargeProd')
+                    internalNonbonded.addPerBondParameter('sigma')
+                    internalNonbonded.addPerBondParameter('epsilon')
+                    numParticles = system.getNumParticles()
+                    atomCharge = [0]*numParticles
+                    atomSigma = [0]*numParticles
+                    atomEpsilon = [0]*numParticles
+                    for i in range(numParticles):
+                        charge, sigma, epsilon = force.getParticleParameters(i)
+                        atomCharge[i] = charge
+                        atomSigma[i] = sigma
+                        atomEpsilon[i] = epsilon
+                    exceptions = {}
+                    for i in range(force.getNumExceptions()):
+                        p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(i)
+                        exceptions[(p1, p2)] = (chargeProd, sigma, epsilon)
+                    for p1 in atomList:
+                        for p2 in atomList:
+                            if p1 == p2:
+                                break
+                            if (p1, p2) in exceptions:
+                                chargeProd, sigma, epsilon = exceptions[(p1, p2)]
+                            elif (p2, p1) in exceptions:
+                                chargeProd, sigma, epsilon = exceptions[(p2, p1)]
+                            else:
+                                chargeProd = atomCharge[p1]*atomCharge[p2]
+                                sigma = 0.5*(atomSigma[p1]+atomSigma[p2])
+                                epsilon = unit.sqrt(atomEpsilon[p1]*atomEpsilon[p2])
+                            if chargeProd._value != 0 or epsilon._value != 0:
+                                internalNonbonded.addBond(p1, p2, [chargeProd, sigma, epsilon])
+                    if internalNonbonded.getNumBonds() > 0:
+                        name = f'mmForce{len(mmVarNames)+1}'
+                        cv.addCollectiveVariable(name, internalNonbonded)
+                        mmVarNames.append(name)
+
+            # Configure the CustomCVForce so lambda_interpolate interpolates between the conventional and ML potentials.
+
+            mlSum = '+'.join(mlVarNames) if len(mlVarNames) > 0 else '0'
+            mmSum = '+'.join(mmVarNames) if len(mmVarNames) > 0 else '0'
+            cv.setEnergyFunction(f'lambda_interpolate*({mlSum}) + (1-lambda_interpolate)*({mmSum})')
+            newSystem.addForce(cv)
+        return newSystem
+
+    def _removeBonds(self, system: openmm.System, atoms: Iterable[int], removeInSet: bool, removeConstraints: bool) -> openmm.System:
+        """Copy a System, removing all bonded interactions between atoms in (or not in) a particular set.
+
+        Parameters
+        ----------
+        system: System
+            the System to copy
+        atoms: Iterable[int]
+            a set of atom indices
+        removeInSet: bool
+            if True, any bonded term connecting atoms in the specified set is removed.  If False,
+            any term that does *not* connect atoms in the specified set is removed
+        removeConstraints: bool
+            if True, remove constraints between pairs of atoms in the set
+
+        Returns
+        -------
+        a newly created System object in which the specified bonded interactions have been removed
+        """
+        atomSet = set(atoms)
+
+        # Create an XML representation of the System.
+
+        import xml.etree.ElementTree as ET
+        xml = openmm.XmlSerializer.serialize(system)
+        root = ET.fromstring(xml)
+
+        # This function decides whether a bonded interaction should be removed.
+
+        def shouldRemove(termAtoms):
+            return all(a in atomSet for a in termAtoms) == removeInSet
+
+        # Remove bonds, angles, and torsions.
+
+        for bonds in root.findall('./Forces/Force/Bonds'):
+            for bond in bonds.findall('Bond'):
+                bondAtoms = [int(bond.attrib[p]) for p in ('p1', 'p2')]
+                if shouldRemove(bondAtoms):
+                    bonds.remove(bond)
+        for angles in root.findall('./Forces/Force/Angles'):
+            for angle in angles.findall('Angle'):
+                angleAtoms = [int(angle.attrib[p]) for p in ('p1', 'p2', 'p3')]
+                if shouldRemove(angleAtoms):
+                    angles.remove(angle)
+        for torsions in root.findall('./Forces/Force/Torsions'):
+            for torsion in torsions.findall('Torsion'):
+                torsionLabels =  ('p1', 'p2', 'p3', 'p4') if 'p1' in torsion.attrib else ('a1', 'a2', 'a3', 'a4', 'b1', 'b2', 'b3', 'b4')
+                torsionAtoms = [int(torsion.attrib[p]) for p in torsionLabels]
+                if shouldRemove(torsionAtoms):
+                    torsions.remove(torsion)
+
+        # Optionally remove constraints.
+
+        if removeConstraints:
+            for constraints in root.findall('./Constraints'):
+                for constraint in constraints.findall('Constraint'):
+                    constraintAtoms = [int(constraint.attrib[p]) for p in ('p1', 'p2')]
+                    if shouldRemove(constraintAtoms):
+                        constraints.remove(constraint)
+
+        # Create a new System from it.
+
+        return openmm.XmlSerializer.deserialize(ET.tostring(root, encoding='unicode'))
+
+    @staticmethod
+    def registerImplFactory(name: str, factory: MLPotentialImplFactory):
+        """Register a new potential function that can be used with MLPotential.
+
+        Parameters
+        ----------
+        name: str
+            the name of the potential function that will be passed to the MLPotential constructor
+        factory: MLPotentialImplFactory
+            a factory object that will be used to create MLPotentialImpl objects
+        """
+        MLPotential._implFactories[name] = factory
+
+
+# Register any potential functions defined by entry points.
+
+for potential in entry_points(group='openmmml.potentials'):
+    MLPotential.registerImplFactory(potential.name, potential.load()())

openmmml-1.6/openmmml/models/__init__.py

@@ -0,0 +1 @@
+from . import anipotential, macepotential, nequippotential, deepmdpotential, aimnet2potential