diff --git a/docs/releasehistory.md b/docs/releasehistory.md index a62cb731..28ce5a75 100644 --- a/docs/releasehistory.md +++ b/docs/releasehistory.md @@ -27,6 +27,7 @@ Please note that all releases prior to a version 1.0.0 are considered pre-releas ### Performance improvements * #1097 Migrates version handling to `versioningit`, which should result in shorter import times. +* #1122 Improves performance of Interchange.from_smirnoff on polymers. ### Documentation improvements diff --git a/openff/interchange/components/toolkit.py b/openff/interchange/components/toolkit.py index 36159755..12629795 100644 --- a/openff/interchange/components/toolkit.py +++ b/openff/interchange/components/toolkit.py @@ -1,19 +1,35 @@ """Utilities for processing and interfacing with the OpenFF Toolkit.""" +from functools import lru_cache from typing import TYPE_CHECKING, Union import networkx import numpy from openff.toolkit import ForceField, Molecule, Quantity, Topology from openff.toolkit.topology._mm_molecule import _SimpleMolecule -from openff.toolkit.typing.engines.smirnoff.parameters import VirtualSiteHandler +from openff.toolkit.typing.engines.smirnoff.parameters import ParameterHandler, VirtualSiteHandler from openff.toolkit.utils.collections import ValidatedList from openff.utilities.utilities import has_package +from openff.interchange.models import ( + PotentialKey, +) + if has_package("openmm") or TYPE_CHECKING: import openmm.app +_IDIVF_1 = Quantity(1.0, "dimensionless") +_PERIODICITIES = { + 1: Quantity(1, "dimensionless"), + 2: Quantity(2, "dimensionless"), + 3: Quantity(3, "dimensionless"), + 4: Quantity(4, "dimensionless"), + 5: Quantity(5, "dimensionless"), + 6: Quantity(6, "dimensionless"), +} + + def _get_num_h_bonds(topology: "Topology") -> int: """Get the number of (covalent) bonds containing a hydrogen atom.""" n_bonds_containing_hydrogen = 0 @@ -202,3 +218,42 @@ def _lookup_virtual_site_parameter( raise ValueError( f"No VirtualSiteType found with {smirks=}, name={name=}, and match={match=}.", ) + + +@lru_cache +def _cache_angle_parameter_lookup( + potential_key: PotentialKey, + parameter_handler: ParameterHandler, +) -> dict[str, Quantity]: + parameter = parameter_handler.parameters[potential_key.id] + + return {parameter_name: getattr(parameter, parameter_name) for parameter_name in ["k", "angle"]} + + +@lru_cache +def _cache_torsion_parameter_lookup( + potential_key: PotentialKey, + parameter_handler: ParameterHandler, + idivf: float | None = None, +) -> dict[str, Quantity]: + smirks = potential_key.id + n = potential_key.mult + parameter = parameter_handler.parameters[smirks] + + if idivf is not None: + # case of non-standard default_idivf in impropers + _idivf = idivf + elif parameter.idivf is None: + # This appears to only come from imports + _idivf = _IDIVF_1 + elif parameter.idivf[n] == 1.0: + _idivf = _IDIVF_1 + else: + _idivf = Quantity(parameter.idivf[n], "dimensionless") + + return { + "k": parameter.k[n], + "periodicity": _PERIODICITIES[parameter.periodicity[n]], + "phase": parameter.phase[n], + "idivf": _idivf, + } diff --git a/openff/interchange/interop/openmm/_import/_import.py b/openff/interchange/interop/openmm/_import/_import.py index 2a0a6389..c8604e48 100644 --- a/openff/interchange/interop/openmm/_import/_import.py +++ b/openff/interchange/interop/openmm/_import/_import.py @@ -135,8 +135,6 @@ def from_openmm( def _convert_constraints( system: "openmm.System", ) -> ConstraintCollection | None: - from openff.toolkit import unit - from openff.interchange.components.potentials import Potential from openff.interchange.models import BondKey, PotentialKey @@ -164,7 +162,7 @@ def _convert_constraints( potential_key = PotentialKey(id=f"Constraint{index}") _keys[distance] = potential_key constraints.potentials[potential_key] = Potential( - parameters={"distance": distance * unit.nanometer}, + parameters={"distance": Quantity(distance, "nanometer")}, ) for index in range(system.getNumConstraints()): diff --git a/openff/interchange/smirnoff/_nonbonded.py b/openff/interchange/smirnoff/_nonbonded.py index 453f60c0..38f4f3da 100644 --- a/openff/interchange/smirnoff/_nonbonded.py +++ b/openff/interchange/smirnoff/_nonbonded.py @@ -53,6 +53,17 @@ LibraryChargeHandler, ] +_ZERO_CHARGE = Quantity(0.0, unit.elementary_charge) + + +@functools.lru_cache(None) +def _add_charges( + charge1: float, + charge2: float, +) -> "Quantity": + """Add two charges together.""" + return Quantity(charge1 + charge2, "elementary_charge") + def _upconvert_vdw_handler(vdw_handler: vdWHandler): """Given a vdW with version 0.3 or 0.4, up-convert to 0.4 or short-circuit if already 0.4.""" @@ -362,7 +373,6 @@ def _get_charges( # the "charge" and "charge_increment" keys may not appear in that order, so # we "add" the charge whether or not the increment was already applied. # There should be a better way to do this. - charges[atom_index] = charges.get(atom_index, 0.0) + parameter_value.m else: @@ -781,7 +791,6 @@ def _find_reference_matches( @classmethod def _assign_charges_from_molecules( cls, - topology: Topology, unique_molecule: Molecule, molecules_with_preset_charges=list[Molecule] | None, ) -> tuple[bool, dict, dict]: @@ -851,7 +860,6 @@ def store_matches( unique_molecule = topology.molecule(unique_molecule_index) flag, matches, potentials = self._assign_charges_from_molecules( - topology, unique_molecule, molecules_with_preset_charges, ) diff --git a/openff/interchange/smirnoff/_valence.py b/openff/interchange/smirnoff/_valence.py index 63d5ae08..b97318d4 100644 --- a/openff/interchange/smirnoff/_valence.py +++ b/openff/interchange/smirnoff/_valence.py @@ -20,6 +20,7 @@ ProperTorsionCollection, ) from openff.interchange.components.potentials import Potential, WrappedPotential +from openff.interchange.components.toolkit import _cache_angle_parameter_lookup, _cache_torsion_parameter_lookup from openff.interchange.exceptions import ( DuplicateMoleculeError, InvalidParameterHandlerError, @@ -448,15 +449,16 @@ def store_potentials(self, parameter_handler: AngleHandler) -> None: """ for potential_key in self.key_map.values(): - smirks = potential_key.id - parameter = parameter_handler.parameters[smirks] - potential = Potential( - parameters={ - parameter_name: getattr(parameter, parameter_name) - for parameter_name in self.potential_parameters() + self.potentials.update( + { + potential_key: Potential( + parameters=_cache_angle_parameter_lookup( + potential_key, + parameter_handler, + ), + ), }, ) - self.potentials[potential_key] = potential class SMIRNOFFProperTorsionCollection(SMIRNOFFCollection, ProperTorsionCollection): @@ -550,11 +552,11 @@ def store_potentials(self, parameter_handler: ProperTorsionHandler) -> None: """ for topology_key, potential_key in self.key_map.items(): - smirks = potential_key.id - n = potential_key.mult - parameter = parameter_handler.parameters[smirks] - if topology_key.bond_order: # type: ignore[union-attr] + smirks = potential_key.id + n = potential_key.mult + parameter = parameter_handler.parameters[smirks] + bond_order = topology_key.bond_order # type: ignore[union-attr] data = parameter.k_bondorder[n] coeffs = _get_interpolation_coeffs( @@ -580,12 +582,7 @@ def store_potentials(self, parameter_handler: ProperTorsionHandler) -> None: {pot: coeff for pot, coeff in zip(pots, coeffs)}, ) else: - parameters = { - "k": parameter.k[n], - "periodicity": parameter.periodicity[n] * unit.dimensionless, - "phase": parameter.phase[n], - "idivf": parameter.idivf[n] * unit.dimensionless, - } + parameters = _cache_torsion_parameter_lookup(potential_key, parameter_handler) potential = Potential(parameters=parameters) # type: ignore[assignment] self.potentials[potential_key] = potential @@ -734,11 +731,11 @@ def store_potentials(self, parameter_handler: ImproperTorsionHandler) -> None: # Assumed to be a numerical value idivf = _default_idivf * unit.dimensionless - parameters = { - "k": parameter.k[n], - "periodicity": parameter.periodicity[n] * unit.dimensionless, - "phase": parameter.phase[n], - "idivf": idivf, - } - potential = Potential(parameters=parameters) - self.potentials[potential_key] = potential + # parameter keys happen to be the same as keys in proper torsions + self.potentials[potential_key] = Potential( + parameters=_cache_torsion_parameter_lookup( + potential_key, + parameter_handler, + idivf=idivf, + ), + )