diff --git a/openmmml/models/aimnet2potential.py b/openmmml/models/aimnet2potential.py index 186d15a..40dbe2f 100644 --- a/openmmml/models/aimnet2potential.py +++ b/openmmml/models/aimnet2potential.py @@ -134,11 +134,8 @@ 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) @@ -146,10 +143,12 @@ def addForces(self, # 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 diff --git a/openmmml/models/anipotential.py b/openmmml/models/anipotential.py index 0347647..1732506 100644 --- a/openmmml/models/anipotential.py +++ b/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 diff --git a/openmmml/models/asepotential.py b/openmmml/models/asepotential.py index b0c9fa8..94a9f89 100644 --- a/openmmml/models/asepotential.py +++ b/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) diff --git a/openmmml/models/fennixpotential.py b/openmmml/models/fennixpotential.py index 885fb49..2f22d54 100644 --- a/openmmml/models/fennixpotential.py +++ b/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) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index 742323e..f90b9ae 100644 --- a/openmmml/models/macepotential.py +++ b/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 diff --git a/openmmml/models/nequippotential.py b/openmmml/models/nequippotential.py index ebb1811..ebdd6be 100644 --- a/openmmml/models/nequippotential.py +++ b/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 diff --git a/openmmml/models/orbpotential.py b/openmmml/models/orbpotential.py index b9584a9..91e8d60 100644 --- a/openmmml/models/orbpotential.py +++ b/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,23 +105,21 @@ 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)) @@ -131,8 +127,4 @@ def _computeOrb(state, atoms, indices, periodic, device, model, adapter, conserv 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)) diff --git a/openmmml/models/torchmdnetpotential.py b/openmmml/models/torchmdnetpotential.py index daecfd5..d6fb4d8 100644 --- a/openmmml/models/torchmdnetpotential.py +++ b/openmmml/models/torchmdnetpotential.py @@ -191,10 +191,6 @@ def addForces(self, batch = torch.zeros_like(numbers, requires_grad=False) else: batch = torch.tensor(batch, dtype=torch.long, device=device, requires_grad=False) - if atoms is None: - indices = None - else: - indices = np.array(atoms) # Create the PythonForce and add it to the System. @@ -204,11 +200,12 @@ def addForces(self, batch=batch, lengthScale=self.lengthScale, energyScale=self.energyScale, - 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: @@ -218,7 +215,7 @@ def getMLLongRange(self) -> bool | None: return None class _ComputeTorchMDNet(object): - def __init__(self, model, numbers, charge, batch, lengthScale, energyScale, indices, periodic): + def __init__(self, model, numbers, charge, batch, lengthScale, energyScale, periodic): self.model = model self.compiled_model = None self.numbers = numbers @@ -226,17 +223,13 @@ def __init__(self, model, numbers, charge, batch, lengthScale, energyScale, indi self.batch = batch self.lengthScale = lengthScale self.energyScale = energyScale - self.indices = indices self.periodic = periodic self.has_recompiled = False def __call__(self, state): import torch positions = state.getPositions(asNumpy=True).value_in_unit(unit.nanometer) - numAtoms = positions.shape[0] positions = torch.tensor(positions, dtype=torch.float32, device=self.numbers.device) - if self.indices is not None: - positions = positions[self.indices] positions.requires_grad_(True) if self.periodic: cell = torch.tensor(state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.nanometer), dtype=torch.float32, device=self.numbers.device)/self.lengthScale @@ -261,8 +254,4 @@ def __call__(self, state): energy = self.compiled_model(z=self.numbers, pos=positions/self.lengthScale, batch=self.batch, q=self.charge, box=cell)[0]*self.energyScale energy.backward() forces = (-positions.grad).detach().cpu().numpy() - if self.indices is not None: - f = np.zeros((numAtoms, 3), dtype=np.float32) - f[self.indices] = forces - forces = f return energy, forces diff --git a/setup.py b/setup.py index 5ebcb11..0d3a9c8 100644 --- a/setup.py +++ b/setup.py @@ -38,7 +38,7 @@ classifiers=CLASSIFIERS.splitlines(), packages=find_packages(), zip_safe=False, - install_requires=['numpy', 'openmm >= 8.5.2'], + install_requires=['numpy', 'openmm >= 8.6.1'], entry_points={ 'openmmml.potentials': [ # Custom AIMNet2 model supplied by path.