From 396d472dfd0ae8b89438b7fb2b33c21c7fb5c9ef Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Thu, 21 Nov 2024 11:11:37 -0600 Subject: [PATCH 1/2] FIX: More consistently use existing charge caching --- .github/workflows/ci.yaml | 4 +- examples/protein_ligand/protein_ligand.ipynb | 6 +-- openff/interchange/common/_nonbonded.py | 5 +- .../interop/amber/export/_export.py | 3 +- .../interop/lammps/export/export.py | 2 +- .../interop/openmm/_import/_import.py | 9 ++-- .../interchange/interop/openmm/_nonbonded.py | 6 ++- openff/interchange/smirnoff/_gromacs.py | 10 ++-- openff/interchange/smirnoff/_nonbonded.py | 46 ++++--------------- 9 files changed, 25 insertions(+), 66 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 4ae4eb072..c44a103bb 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -95,9 +95,7 @@ jobs: run: micromamba install "foyer >=0.12.1" -c conda-forge -yq - name: Run tests - if: always() - run: | - python -m pytest $COV openff/interchange/ -r fExs -n logical --durations=10 + run: python -m pytest $COV openff/interchange/ -r fExs -n logical --durations=10 - name: Run small molecule regression tests if: ${{ matrix.openeye == true && matrix.openmm == true }} diff --git a/examples/protein_ligand/protein_ligand.ipynb b/examples/protein_ligand/protein_ligand.ipynb index 6ac020ea5..059eb0375 100644 --- a/examples/protein_ligand/protein_ligand.ipynb +++ b/examples/protein_ligand/protein_ligand.ipynb @@ -554,9 +554,7 @@ "id": "47", "metadata": {}, "source": [ - "### GROMACS\n", - "\n", - "Interchange's GROMACS exporter is a little slow for biopolymers; this will be faster in a future release." + "### GROMACS" ] }, { @@ -717,7 +715,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.4" + "version": "3.10.15" }, "vscode": { "interpreter": { diff --git a/openff/interchange/common/_nonbonded.py b/openff/interchange/common/_nonbonded.py index 551c157c8..829bfe2d6 100644 --- a/openff/interchange/common/_nonbonded.py +++ b/openff/interchange/common/_nonbonded.py @@ -7,7 +7,6 @@ from openff.interchange._annotations import _DistanceQuantity, _ElementaryChargeQuantity from openff.interchange.components.potentials import Collection -from openff.interchange.constants import _PME from openff.interchange.models import ( LibraryChargeTopologyKey, TopologyKey, @@ -95,9 +94,7 @@ class ElectrostaticsCollection(_NonbondedCollection): "Ewald3D-ConductingBoundary", "cutoff", "no-cutoff", - ] = Field( - _PME, - ) # type: ignore[assignment] + ] = Field("Ewald3D-ConductingBoundary") nonperiodic_potential: Literal["Coulomb", "cutoff", "no-cutoff"] = Field("Coulomb") exception_potential: Literal["Coulomb"] = Field("Coulomb") diff --git a/openff/interchange/interop/amber/export/_export.py b/openff/interchange/interop/amber/export/_export.py index 993a0a3bb..5e525a997 100644 --- a/openff/interchange/interop/amber/export/_export.py +++ b/openff/interchange/interop/amber/export/_export.py @@ -486,8 +486,7 @@ def to_prmtop(interchange: "Interchange", file_path: Path | str): prmtop.write("%FLAG CHARGE\n%FORMAT(5E16.8)\n") charges = [ - charge.m_as(unit.e) * AMBER_COULOMBS_CONSTANT - for charge in interchange["Electrostatics"]._get_charges().values() + charge.m_as(unit.e) * AMBER_COULOMBS_CONSTANT for charge in interchange["Electrostatics"].charges.values() ] text_blob = "".join([f"{val:16.8E}" for val in charges]) _write_text_blob(prmtop, text_blob) diff --git a/openff/interchange/interop/lammps/export/export.py b/openff/interchange/interop/lammps/export/export.py index 9f25a50ac..6ba76b7ca 100644 --- a/openff/interchange/interop/lammps/export/export.py +++ b/openff/interchange/interop/lammps/export/export.py @@ -266,7 +266,7 @@ def _write_atoms(lmp_file: IO, interchange: Interchange, atom_type_map: dict): vdw_handler = interchange["vdW"] - charges = interchange["Electrostatics"]._get_charges() + charges = interchange["Electrostatics"].charges positions = interchange.positions.m_as(unit.angstrom) # type: ignore[union-attr] """ diff --git a/openff/interchange/interop/openmm/_import/_import.py b/openff/interchange/interop/openmm/_import/_import.py index fd4fd7686..b90bf1dcd 100644 --- a/openff/interchange/interop/openmm/_import/_import.py +++ b/openff/interchange/interop/openmm/_import/_import.py @@ -4,7 +4,7 @@ from openff.toolkit import Quantity, Topology from openff.utilities.utilities import has_package, requires_package -from openff.interchange.common._nonbonded import vdWCollection +from openff.interchange.common._nonbonded import ElectrostaticsCollection, vdWCollection from openff.interchange.common._valence import ( AngleCollection, BondCollection, @@ -12,9 +12,6 @@ ProperTorsionCollection, ) from openff.interchange.exceptions import UnsupportedImportError -from openff.interchange.interop.openmm._import._nonbonded import ( - BasicElectrostaticsCollection, -) from openff.interchange.interop.openmm._import.compat import _check_compatible_inputs from openff.interchange.warnings import MissingPositionsWarning @@ -177,7 +174,7 @@ def _convert_constraints( def _convert_nonbonded_force( force: "openmm.NonbondedForce", -) -> tuple[vdWCollection, BasicElectrostaticsCollection]: +) -> tuple[vdWCollection, ElectrostaticsCollection]: from openff.units.openmm import from_openmm as from_openmm_quantity from openff.interchange.components.potentials import Potential @@ -189,7 +186,7 @@ def _convert_nonbonded_force( ) vdw = vdWCollection() - electrostatics = BasicElectrostaticsCollection(version=0.4, scale_14=0.833333) + electrostatics = ElectrostaticsCollection(version=0.4, scale_14=0.833333) n_parametrized_particles = force.getNumParticles() diff --git a/openff/interchange/interop/openmm/_nonbonded.py b/openff/interchange/interop/openmm/_nonbonded.py index 5f4de68a9..3f24171fd 100644 --- a/openff/interchange/interop/openmm/_nonbonded.py +++ b/openff/interchange/interop/openmm/_nonbonded.py @@ -367,7 +367,7 @@ def _create_single_nonbonded_force( ) if data.electrostatics_collection is not None: - partial_charges = data.electrostatics_collection._get_charges() + partial_charges = data.electrostatics_collection.charges # mapping between (openmm) index of each atom and the (openmm) index of each virtual particle # of that parent atom (if any) @@ -394,7 +394,9 @@ def _create_single_nonbonded_force( other_top_key = SingleAtomChargeTopologyKey( this_atom_index=atom_index, ) + partial_charge = partial_charges[other_top_key].m_as(unit.e) + else: partial_charge = 0.0 @@ -941,7 +943,7 @@ def _set_particle_parameters( # handling for electrostatics_force = None electrostatics: ElectrostaticsCollection = data.electrostatics_collection - partial_charges = electrostatics._get_charges() + partial_charges = electrostatics.charges vdw: vdWCollection = data.vdw_collection diff --git a/openff/interchange/smirnoff/_gromacs.py b/openff/interchange/smirnoff/_gromacs.py index 27e876fc7..0eab05483 100644 --- a/openff/interchange/smirnoff/_gromacs.py +++ b/openff/interchange/smirnoff/_gromacs.py @@ -110,7 +110,7 @@ def _convert( try: vdw_collection = interchange["vdW"] - electrostatics_collection = interchange["Electrostatics"] + interchange["Electrostatics"] except KeyError: raise UnsupportedExportError("Plugins not implemented.") @@ -146,8 +146,6 @@ def _convert( vdw_parameters = vdw_collection.potentials[vdw_collection.key_map[key]].parameters - charge = electrostatics_collection._get_charges()[key] - # Build atom types system.atom_types[atom_type_name] = LennardJonesAtomType( name=_atom_atom_type_map[atom], @@ -168,8 +166,6 @@ def _convert( vdw_parameters = vdw_collection.potentials[vdw_collection.key_map[virtual_site_key]].parameters - charge = electrostatics_collection._get_charges()[key] - # TODO: Separate class for "atom types" representing virtual sites? system.atom_types[atom_type_name] = LennardJonesAtomType( name=_atom_atom_type_map[virtual_site_key], @@ -185,7 +181,7 @@ def _convert( _partial_charges: dict[int | VirtualSiteKey, float] = dict() # Indexed by particle (atom or virtual site) indices - for key, charge in interchange["Electrostatics"]._get_charges().items(): + for key, charge in interchange["Electrostatics"].charges.items(): if type(key) is TopologyKey: _partial_charges[key.atom_indices[0]] = charge elif type(key) is VirtualSiteKey: @@ -585,7 +581,7 @@ def _convert_virtual_sites( residue_index=molecule.atoms[0].residue_index, residue_name=molecule.atoms[0].residue_name, charge_group_number=1, - charge=interchange["Electrostatics"]._get_charges()[virtual_site_key], + charge=interchange["Electrostatics"].charges[virtual_site_key], mass=Quantity(0.0, unit.dalton), ), ) diff --git a/openff/interchange/smirnoff/_nonbonded.py b/openff/interchange/smirnoff/_nonbonded.py index de4b800aa..2c81db3d5 100644 --- a/openff/interchange/smirnoff/_nonbonded.py +++ b/openff/interchange/smirnoff/_nonbonded.py @@ -52,21 +52,6 @@ LibraryChargeHandler, ] -_ZERO_CHARGE = Quantity(0.0, unit.elementary_charge) - - -@unit.wraps( - ret=unit.elementary_charge, - args=(unit.elementary_charge, unit.elementary_charge), - strict=True, -) -def _add_charges( - charge1: "Quantity", - charge2: "Quantity", -) -> "Quantity": - """Add two charges together.""" - return charge1 + charge2 - 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.""" @@ -331,6 +316,7 @@ def _get_charges( ) -> dict[TopologyKey | LibraryChargeTopologyKey | VirtualSiteKey, _ElementaryChargeQuantity]: """Get the total partial charge on each atom or particle.""" # Keyed by index for atoms and by VirtualSiteKey for virtual sites. + # deal with untless (float, implicit elementary_charge) values until returning charges: dict[VirtualSiteKey | int, _ElementaryChargeQuantity] = dict() for topology_key, potential_key in self.key_map.items(): @@ -344,19 +330,14 @@ def _get_charges( "virtual sites, not by a `ChargeIncrementModelHandler`.", ) - total_charge: Quantity = numpy.sum(parameter_value) # assumes virtual sites can only have charges determined in one step - - charges[topology_key] = -1.0 * total_charge + charges[topology_key] = -1.0 * numpy.sum(parameter_value) # Apply increments to "orientation" atoms for i, increment in enumerate(parameter_value): orientation_atom_index = topology_key.orientation_atom_indices[i] - charges[orientation_atom_index] = _add_charges( - charges.get(orientation_atom_index, _ZERO_CHARGE), - increment, - ) + charges[orientation_atom_index] = charges.get(orientation_atom_index, 0.0) + increment.m elif parameter_key == "charge": assert len(topology_key.atom_indices) == 1 @@ -369,10 +350,7 @@ def _get_charges( "molecules_with_preset_charges", "ExternalSource", ): - charges[atom_index] = _add_charges( - charges.get(atom_index, _ZERO_CHARGE), - parameter_value, - ) + charges[atom_index] = charges.get(atom_index, 0.0) + parameter_value.m elif potential_key.associated_handler in ( # type: ignore[operator] "ChargeIncrementModelHandler" @@ -381,10 +359,7 @@ def _get_charges( # we "add" the charge whether or not the increment was already applied. # There should be a better way to do this. - charges[atom_index] = _add_charges( - charges.get(atom_index, _ZERO_CHARGE), - parameter_value, - ) + charges[atom_index] = charges.get(atom_index, 0.0) + parameter_value.m else: raise RuntimeError( @@ -396,13 +371,10 @@ def _get_charges( atom_index = topology_key.atom_indices[0] - charges[atom_index] = _add_charges( - charges.get(atom_index, _ZERO_CHARGE), - parameter_value, - ) + charges[atom_index] = charges.get(atom_index, 0.0) + parameter_value.m logger.info( - "Charge section ChargeIncrementModel, applying charge increment from atom " + "Charge section ChargeIncrementModel, applying charge increment from atom " # type: ignore[union-attr] f"{topology_key.this_atom_index} to atoms {topology_key.other_atom_indices}", ) @@ -424,7 +396,7 @@ def _get_charges( if include_virtual_sites: returned_charges[index] = charge - return returned_charges + return {key: Quantity(val, "elementary_charge") for key, val in returned_charges.items()} @classmethod def parameter_handler_precedence(cls) -> list[str]: @@ -952,7 +924,7 @@ def store_matches( topology_charges = [0.0] * topology.n_atoms - for key, val in self._get_charges().items(): + for key, val in self.charges.items(): topology_charges[key.atom_indices[0]] = val.m # TODO: Better data structures in Topology.identical_molecule_groups will make this From 9f03ef2ae9fe855949c7223e4b11c20dcdbf3108 Mon Sep 17 00:00:00 2001 From: Matt Thompson Date: Wed, 15 Jan 2025 08:18:48 -0600 Subject: [PATCH 2/2] Update openff/interchange/smirnoff/_nonbonded.py Co-authored-by: Josh A. Mitchell --- openff/interchange/smirnoff/_nonbonded.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openff/interchange/smirnoff/_nonbonded.py b/openff/interchange/smirnoff/_nonbonded.py index 2c81db3d5..4c22963ec 100644 --- a/openff/interchange/smirnoff/_nonbonded.py +++ b/openff/interchange/smirnoff/_nonbonded.py @@ -316,7 +316,7 @@ def _get_charges( ) -> dict[TopologyKey | LibraryChargeTopologyKey | VirtualSiteKey, _ElementaryChargeQuantity]: """Get the total partial charge on each atom or particle.""" # Keyed by index for atoms and by VirtualSiteKey for virtual sites. - # deal with untless (float, implicit elementary_charge) values until returning + # work in unitless (float, implicit elementary_charge) values until returning charges: dict[VirtualSiteKey | int, _ElementaryChargeQuantity] = dict() for topology_key, potential_key in self.key_map.items():