Use setParticles() to restrict atoms

openmmml/models/aimnet2potential.py

@@ -134,22 +134,21 @@ def addForces(self,
         # Create the PyTorch model that will be invoked by OpenMM.
 
         includedAtoms = list(topology.atoms())
-        if atoms is None:
-            indices = None
-        else:
+        if atoms is not None:
             includedAtoms = [includedAtoms[i] for i in atoms]
-            indices = np.array(atoms)
         numbers = torch.tensor([[atom.element.atomic_number for atom in includedAtoms]], device=device)
         charge = torch.tensor([args.get('charge', 0)], dtype=torch.float32, device=device)
         multiplicity = torch.tensor([args.get('multiplicity', 1)], dtype=torch.float32, device=device)
         periodic = topology.getPeriodicBoxVectors() is not None
 
         # Create the PythonForce and add it to the System.
 
-        compute = partial(_computeAIMNet2, model=model, numbers=numbers, charge=charge, multiplicity=multiplicity, indices=indices, periodic=periodic)
+        compute = partial(_computeAIMNet2, model=model, numbers=numbers, charge=charge, multiplicity=multiplicity, periodic=periodic)
         force = openmm.PythonForce(compute)
         force.setForceGroup(forceGroup)
         force.setUsesPeriodicBoundaryConditions(periodic)
+        if atoms is not None:
+            force.setParticles(atoms)
         system.addForce(force)
 
     def getMLLongRange(self) -> bool | None:
@@ -163,12 +162,9 @@ def getMLLongRange(self) -> bool | None:
         # supported model has different behavior, this must be updated.
         return False
 
-def _computeAIMNet2(state, model, numbers, charge, multiplicity, indices, periodic):
+def _computeAIMNet2(state, model, numbers, charge, multiplicity, periodic):
     import torch
     positions = torch.tensor(state.getPositions(asNumpy=True).value_in_unit(unit.angstrom), dtype=torch.float32, device=numbers.device)
-    numAtoms = positions.shape[0]
-    if indices is not None:
-        positions = positions[indices]
     args = {'coord': positions.unsqueeze(0),
             'numbers': numbers,
             'charge': charge,
@@ -180,8 +176,4 @@ def _computeAIMNet2(state, model, numbers, charge, multiplicity, indices, period
     energyScale = (unit.ev/unit.item).conversion_factor_to(unit.kilojoules_per_mole)
     energy = float(energyScale*result["energy"].sum().detach())
     forces = (10.0*energyScale*result["forces"]).detach().cpu().numpy()[0]
-    if indices is not None:
-        f = np.zeros((numAtoms, 3), dtype=np.float32)
-        f[indices] = forces
-        forces = f
     return energy, forces

openmmml/models/anipotential.py

@@ -93,10 +93,6 @@ def addForces(self,
         if atoms is not None:
             includedAtoms = [includedAtoms[i] for i in atoms]
         species = torch.tensor([[atom.element.atomic_number for atom in includedAtoms]], device=device)
-        if atoms is None:
-            indices = None
-        else:
-            indices = np.array(atoms)
         periodic = topology.getPeriodicBoxVectors() is not None or system.usesPeriodicBoundaryConditions()
         if periodic:
             pbc = torch.tensor([True, True, True], dtype=torch.bool, device=device)
@@ -108,24 +104,21 @@ def addForces(self,
         compute = partial(_computeANI,
                           model=model,
                           species=species,
-                          pbc=pbc,
-                          indices=indices)
+                          pbc=pbc)
         force = openmm.PythonForce(compute)
         force.setForceGroup(forceGroup)
         force.setUsesPeriodicBoundaryConditions(periodic)
+        if atoms is not None:
+            force.setParticles(atoms)
         system.addForce(force)
 
     def getMLLongRange(self) -> bool | None:
         return False
 
-def _computeANI(state, model, species, pbc, indices):
+def _computeANI(state, model, species, pbc):
     import torch
-    import numpy as np
     import torchani
     positions = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom)
-    numAtoms = positions.shape[0]
-    if indices is not None:
-        positions = positions[indices]
     positions = torch.tensor(positions, dtype=torch.float32, device=species.device)
     if pbc is None:
         boxvectors = None
@@ -140,8 +133,4 @@ def _computeANI(state, model, species, pbc, indices):
     energy *= torchani.units.hartree2kjoulemol(1)
     energy.backward()
     forces = (-positions.grad[0]).detach().cpu().numpy()
-    if indices is not None:
-        f = np.zeros((numAtoms, 3), dtype=np.float32)
-        f[indices] = forces
-        forces = f
     return energy, forces

openmmml/models/asepotential.py

@@ -82,11 +82,8 @@ def addForces(self,
         if any(atom.element is None for atom in topology.atoms()):
             raise ValueError('All atoms in the Topology must have elements defined.')
         includedAtoms = list(topology.atoms())
-        if atoms is None:
-            indices = None
-        else:
+        if atoms is not None:
             includedAtoms = [includedAtoms[i] for i in atoms]
-            indices = np.array(atoms)
         if 'aseAtoms' in args:
             # The user provided an Atoms object.
 
@@ -109,26 +106,21 @@ def addForces(self,
 
         # Create the PythonForce and add it to the System.
 
-        compute = partial(_computeASE, atoms=aseAtoms, indices=indices)
+        compute = partial(_computeASE, atoms=aseAtoms)
         force = openmm.PythonForce(compute)
         force.setForceGroup(forceGroup)
         force.setUsesPeriodicBoundaryConditions(any(aseAtoms.get_pbc()))
+        if atoms is not None:
+            force.setParticles(atoms)
         system.addForce(force)
 
 
-def _computeASE(state, atoms, indices):
+def _computeASE(state, atoms):
     import ase.units
     positions = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom)
-    numAtoms = positions.shape[0]
-    if indices is not None:
-        positions = positions[indices]
     atoms.set_positions(positions)
     if any(atoms.get_pbc()):
         atoms.set_cell(state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom))
     energy = atoms.get_potential_energy(apply_constraint=False)
     forces = atoms.get_forces(apply_constraint=False)
-    if indices is not None:
-        f = np.zeros((numAtoms, 3), dtype=np.float32)
-        f[indices] = forces
-        forces = f
     return energy/(ase.units.kJ/ase.units.mol), forces*10/(ase.units.kJ/ase.units.mol)

openmmml/models/fennixpotential.py

@@ -135,10 +135,8 @@ def addForces(self,
 
         # Get the atoms that should be included.
         includedAtoms = list(topology.atoms())
-        indices = None
         if atoms is not None:
             includedAtoms = [includedAtoms[i] for i in atoms]
-            indices = np.array(atoms, dtype=int)
 
         # Prepare inputs to the model that remain constant from step to step.
         species = jnp.array([atom.element.atomic_number for atom in includedAtoms], dtype=jnp.int32)
@@ -151,9 +149,11 @@ def addForces(self,
 
         # Create the PythonForce and add it to the System.
         periodic = (topology.getPeriodicBoxVectors() is not None) or system.usesPeriodicBoundaryConditions()
-        force = openmm.PythonForce(_ComputeFeNNix(model, energyScale, forceScale, indices, inputs, periodic, useDouble))
+        force = openmm.PythonForce(_ComputeFeNNix(model, energyScale, forceScale, inputs, periodic, useDouble))
         force.setForceGroup(forceGroup)
         force.setUsesPeriodicBoundaryConditions(periodic)
+        if atoms is not None:
+            force.setParticles(atoms)
         system.addForce(force)
 
     def getMLLongRange(self) -> bool | None:
@@ -164,11 +164,10 @@ def getMLLongRange(self) -> bool | None:
 
 
 class _ComputeFeNNix:
-    def __init__(self, model, energyScale, forceScale, indices, inputs, periodic, useDouble):
+    def __init__(self, model, energyScale, forceScale, inputs, periodic, useDouble):
         self.model = model
         self.energyScale = energyScale
         self.forceScale = forceScale
-        self.indices = indices
         self.inputs = inputs
         self.periodic = periodic
         self.useDouble = useDouble
@@ -179,9 +178,6 @@ def __call__(self, state):
 
         # Load coordinates and box vectors from the state.
         positions = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom)
-        numAtoms = positions.shape[0]
-        if self.indices is not None:
-            positions = positions[self.indices]
         if self.periodic:
             cells = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom).reshape(1, 3, 3)
 
@@ -194,18 +190,14 @@ def __call__(self, state):
             jaxEnergy, jaxForces = modelOutputs[:2]
             energy = jaxEnergy.item() * self.energyScale
             jaxForces *= self.forceScale
-            if self.indices is None:
-                forces = np.asarray(jaxForces)
-            else:
-                forces = np.zeros((numAtoms, 3), dtype=jaxForces.dtype)
-                forces[self.indices] = jaxForces
+            forces = np.asarray(jaxForces)
 
         return energy, forces
 
     def __getstate__(self):
-        return (self.model.to_dict(), self.energyScale, self.forceScale, self.indices, self.inputs, self.periodic, self.useDouble)
+        return (self.model.to_dict(), self.energyScale, self.forceScale, self.inputs, self.periodic, self.useDouble)
 
     def __setstate__(self, pickle_state):
         import fennol
-        model_dict, self.energyScale, self.forceScale, self.indices, self.inputs, self.periodic, self.useDouble = pickle_state
+        model_dict, self.energyScale, self.forceScale, self.inputs, self.periodic, self.useDouble = pickle_state
         self.model = fennol.FENNIX(**model_dict)

openmmml/models/macepotential.py

@@ -214,10 +214,6 @@ def addForces(
             torch.tensor(atomic_numbers_to_indices(atomicNumbers, z_table=zTable), dtype=torch.long, device=device).unsqueeze(-1),
             num_classes=len(zTable))
 
-        if atoms is None:
-            indices = None
-        else:
-            indices = np.array(atoms)
         periodic = (topology.getPeriodicBoxVectors() is not None) or system.usesPeriodicBoundaryConditions()
 
         # Create the PythonForce and add it to the System.
@@ -231,11 +227,12 @@ def addForces(
                           returnEnergyType=returnEnergyType,
                           charge=torch.tensor([float(args.get('charge', 0))], dtype=dtype, device=device, requires_grad=False),
                           multiplicity=torch.tensor([float(args.get('multiplicity', 1))], dtype=dtype, device=device, requires_grad=False),
-                          indices=indices,
                           periodic=periodic)
         force = openmm.PythonForce(compute)
         force.setForceGroup(forceGroup)
         force.setUsesPeriodicBoundaryConditions(periodic)
+        if atoms is not None:
+            force.setParticles(atoms)
         system.addForce(force)
 
     def getMLLongRange(self) -> bool | None:
@@ -245,15 +242,12 @@ def getMLLongRange(self) -> bool | None:
         return None
 
 
-def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, charge, multiplicity, indices, periodic):
+def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, charge, multiplicity, periodic):
     import torch
     from mace.data.neighborhood import get_neighborhood
     energyScale = 96.4853
     lengthScale = 10.0
     positions = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom)
-    numAtoms = positions.shape[0]
-    if indices is not None:
-        positions = positions[indices]
     if periodic:
         cell = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom)
     else:
@@ -276,8 +270,4 @@ def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, ch
     results = model(inputDict, compute_force=True)
     energy = float(results[returnEnergyType].detach())*energyScale
     forces = (results["forces"]*energyScale*lengthScale).detach().cpu().numpy()
-    if indices is not None:
-        f = np.zeros((numAtoms, 3), dtype=(np.float64 if dtype == torch.float64 else np.float32))
-        f[indices] = forces
-        forces = f
     return energy, forces

openmmml/models/nequippotential.py

@@ -199,10 +199,6 @@ def addForces(
         else:
             if len(atomTypes) != len(includedAtoms):
                 raise ValueError("The length of atomTypes must be equal to the number of ML atoms in the system.")
-        if atoms is None:
-            indices = None
-        else:
-            indices = np.array(atoms)
         atomTypes = torch.tensor(atomTypes, dtype=torch.long, requires_grad=False, device=device)
         periodic = (topology.getPeriodicBoxVectors() is not None) or system.usesPeriodicBoundaryConditions()
         pbc = torch.tensor([periodic, periodic, periodic], dtype=torch.bool, requires_grad=False, device=device)
@@ -216,22 +212,20 @@ def addForces(
                           cutoff=cutoff,
                           lengthScale=self.lengthScale,
                           energyScale=self.energyScale,
-                          indices=indices,
                           periodic=periodic,
                           pbc=pbc)
         force = openmm.PythonForce(compute)
         force.setForceGroup(forceGroup)
         force.setUsesPeriodicBoundaryConditions(periodic)
+        if atoms is not None:
+            force.setParticles(atoms)
         system.addForce(force)
 
-def _computeNequIP(state, model, atomTypes, cutoff, lengthScale, energyScale, indices, periodic, pbc):
+def _computeNequIP(state, model, atomTypes, cutoff, lengthScale, energyScale, periodic, pbc):
     import torch
     from nequip.data._nl import compute_neighborlist_
     positions = state.getPositions(asNumpy=True).value_in_unit(unit.nanometer)/lengthScale
-    numAtoms = positions.shape[0]
     positions = torch.tensor(positions, dtype=torch.float64, device=atomTypes.device)
-    if indices is not None:
-        positions = positions[indices]
     inputDict = {
         "pos": positions,
         "atom_types": atomTypes,
@@ -243,8 +237,4 @@ def _computeNequIP(state, model, atomTypes, cutoff, lengthScale, energyScale, in
     out = model(inputDict)
     energy = out["total_energy"] * energyScale
     forces = out["forces"].detach().cpu().numpy()
-    if indices is not None:
-        f = np.zeros((numAtoms, 3), dtype=np.float64)
-        f[indices] = forces
-        forces = f
     return energy, forces*energyScale/lengthScale

openmmml/models/orbpotential.py

@@ -95,10 +95,8 @@ def addForces(self,
 
         # Get the atoms that should be included.
         includedAtoms = list(topology.atoms())
-        indices = None
         if atoms is not None:
             includedAtoms = [includedAtoms[i] for i in atoms]
-            indices = np.array(atoms, dtype=int)
 
         # Set up the ASE Atoms object that will be fed to the model.
         numbers = [atom.element.atomic_number for atom in includedAtoms]
@@ -107,32 +105,26 @@ def addForces(self,
         aseAtoms.info['charge'] = charge
         aseAtoms.info['spin'] = multiplicity
 
-        compute = partial(_computeOrb, atoms=aseAtoms, indices=indices, periodic=periodic, device=device, model=model, adapter=adapter, conservative=conservative)
+        compute = partial(_computeOrb, atoms=aseAtoms, periodic=periodic, device=device, model=model, adapter=adapter, conservative=conservative)
         force = openmm.PythonForce(compute)
         force.setForceGroup(forceGroup)
         force.setUsesPeriodicBoundaryConditions(any(aseAtoms.get_pbc()))
+        if atoms is not None:
+            force.setParticles(atoms)
         system.addForce(force)
 
     def getMLLongRange(self) -> bool | None:
         return False
 
-def _computeOrb(state, atoms, indices, periodic, device, model, adapter, conservative):
+def _computeOrb(state, atoms, periodic, device, model, adapter, conservative):
     import ase.units
-    import numpy as np
 
     positions = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom)
-    numAtoms = positions.shape[0]
-    if indices is not None:
-        positions = positions[indices]
     atoms.set_positions(positions)
     if periodic:
         atoms.set_cell(state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom))
 
     result = model.predict(adapter.from_ase_atoms(atoms, device=device))
     energy = result["energy"].item()
     forces = result[model.grad_forces_name if conservative else "forces"].numpy(force=True)
-    if indices is not None:
-        f = np.zeros((numAtoms, 3), dtype=forces.dtype)
-        f[indices] = forces
-        forces = f
     return energy / (ase.units.kJ / ase.units.mol), forces / (ase.units.kJ / (ase.units.mol * ase.units.nm))