diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 4ae4eb072..c3c2bf3c1 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -22,7 +22,6 @@ jobs: name: Test on ${{ matrix.os }}, Python ${{ matrix.python-version }}, OpenMM ${{ matrix.openmm }}, Pydantic ${{ matrix.pydantic-version }}, OpenEye ${{ matrix.openeye }} runs-on: ${{ matrix.os }} strategy: - fail-fast: false matrix: os: - macos-latest diff --git a/devtools/conda-envs/test_env.yaml b/devtools/conda-envs/test_env.yaml index 581ee7876..58d03fde6 100644 --- a/devtools/conda-envs/test_env.yaml +++ b/devtools/conda-envs/test_env.yaml @@ -17,8 +17,8 @@ dependencies: - openmm =8.1.2 # smirnoff-plugins =2024 # de-forcefields # add back after smirnoff-plugins update - - openff-nagl - - openff-nagl-models + - openff-nagl =0.5 + - openff-nagl-models =0.3 - mbuild ~=0.18 - foyer =1 - gmso ~=0.12 diff --git a/docs/releasehistory.md b/docs/releasehistory.md index 733bc9354..73511c940 100644 --- a/docs/releasehistory.md +++ b/docs/releasehistory.md @@ -15,6 +15,7 @@ Please note that all releases prior to a version 1.0.0 are considered pre-releas ### New features +* #1081 `Interchange.from_openmm` now processes virtual sites, but only `openmm.ThreeParticleAverageSite`s. * #1053 Logs, at the level of `logging.INFO`, how charges are assigned by SMIRNOFF force fields to each atom and virtual site. * #1080 HMR is supported with OpenMM when virtual sites are present. diff --git a/docs/using/edges.md b/docs/using/edges.md index 01ceee0ed..236397abf 100644 --- a/docs/using/edges.md +++ b/docs/using/edges.md @@ -67,6 +67,19 @@ For more, see [issue #1005](https://github.com/openforcefield/openff-interchange Keywords: OpenMM, GROMACS, constraints, bond constraints, rigid water +### Virtual site exclusions re-created with "parents" virtual site exclusion policy + +Non-bonded exclusions involving virtual sites (between virtual sites and heavy atoms or between +virtual sites and virtual sites) are not processed. Instead, they are later re-generated assuming the "parents" exclusion policy as defined in the [SMIRNOFF specification](https://openforcefield.github.io/standards/standards/smirnoff/#virtualsites-virtual-sites-for-off-atom-charges). This should re-create typical exclusions in 4- and 5-site water models but may not be appropriate with highly custom virtual site interactions in larger molecules. + +### Virtual sites from multiple sources cannot be mixed + +Combining systems with virtual sites from multiple sources is not fully-featured. For example, this refers to importing a box of TIP4P-containing solvent from OpenMM with a ligand prepared with SMIRNOFF virtual sites parameters. + +### Virtual sites must be listed after heavy atoms each molecule + +It's assumed that, in each molecule in an OpenMM topology, all heavy atoms are listed before any virtual sites. This includes the case of all virtual sites being listed after all heavy atoms, i.e. not collated into molecules/residues. There are no community standards areound particle ordering, but virtual sites are typically listed after heavy atoms in each molecule or residue. + ## Quirks with GROMACS ### Residue indices must begin at 1 diff --git a/openff/interchange/_tests/conftest.py b/openff/interchange/_tests/conftest.py index 25a2ff70a..061341aa8 100644 --- a/openff/interchange/_tests/conftest.py +++ b/openff/interchange/_tests/conftest.py @@ -13,21 +13,28 @@ @pytest.fixture -def sage(): +def sage() -> ForceField: return ForceField("openff-2.0.0.offxml") @pytest.fixture -def sage_unconstrained(): +def sage_unconstrained() -> ForceField: return ForceField("openff_unconstrained-2.0.0.offxml") @pytest.fixture -def sage_no_switch(sage): +def sage_no_switch(sage) -> ForceField: sage["vdW"].switch_width = Quantity(0.0, "angstrom") return sage +@pytest.fixture +def sage_with_tip4p() -> ForceField: + # re-build off of existing fixtures if this gets implemented + # https://github.com/openforcefield/openff-toolkit/issues/1948 + return ForceField("openff-2.0.0.offxml", "tip4p_fb.offxml") + + @pytest.fixture def sage_with_bond_charge(sage): sage["Bonds"].add_parameter( diff --git a/openff/interchange/_tests/interoperability_tests/components/interchange/test_combine.py b/openff/interchange/_tests/interoperability_tests/components/interchange/test_combine.py index 126c3775c..d1b63cd72 100644 --- a/openff/interchange/_tests/interoperability_tests/components/interchange/test_combine.py +++ b/openff/interchange/_tests/interoperability_tests/components/interchange/test_combine.py @@ -7,7 +7,6 @@ def test_combine_after_from_openmm_with_mainline_openmm_force_field( - monkeypatch, popc, sage, ): @@ -17,8 +16,6 @@ def test_combine_after_from_openmm_with_mainline_openmm_force_field( import openmm.app import openmm.unit - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - ligand = MoleculeWithConformer.from_smiles("c1ccccc1") ligand._conformers[0] += Quantity([3, 3, 3], "angstrom") topology = ligand.to_topology() diff --git a/openff/interchange/_tests/test_issues.py b/openff/interchange/_tests/test_issues.py index 09719cc19..ca68c4f8e 100644 --- a/openff/interchange/_tests/test_issues.py +++ b/openff/interchange/_tests/test_issues.py @@ -78,11 +78,9 @@ def test_issue_1022(pack): @skip_if_missing("openmm") -def test_issue_1031(monkeypatch): +def test_issue_1031(): import openmm.app - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - # just grab some small PDB file from the toolkit, doesn't need to be huge, just # needs to include some relevant atom names openmm_topology = openmm.app.PDBFile( diff --git a/openff/interchange/_tests/unit_tests/components/test_interchange.py b/openff/interchange/_tests/unit_tests/components/test_interchange.py index 66e3c0ba9..cbe4043d0 100644 --- a/openff/interchange/_tests/unit_tests/components/test_interchange.py +++ b/openff/interchange/_tests/unit_tests/components/test_interchange.py @@ -411,9 +411,7 @@ def test_from_gromacs_error(self): Interchange.from_gromacs() @skip_if_missing("openmm") - def test_from_openmm_called(self, monkeypatch, simple_interchange): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - + def test_from_openmm_called(self, simple_interchange): topology = simple_interchange.to_openmm_topology() system = simple_interchange.to_openmm() positions = simple_interchange.positions @@ -435,7 +433,7 @@ def test_to_amber(self, simple_interchange): shell=True, ) - def test_from_gromacs_called(self, monkeypatch, simple_interchange): + def test_from_gromacs_called(self, simple_interchange, monkeypatch): monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") simple_interchange.to_gromacs(prefix="tmp_") diff --git a/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_compat.py b/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_compat.py index 41264c000..a227eca85 100644 --- a/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_compat.py +++ b/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_compat.py @@ -12,9 +12,7 @@ class TestUnsupportedCases: @pytest.mark.filterwarnings("ignore:.*are you sure you don't want to pass positions") - def test_error_topology_mismatch(self, monkeypatch, sage_unconstrained, ethanol): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - + def test_error_topology_mismatch(self, sage_unconstrained, ethanol): topology = ethanol.to_topology() topology.box_vectors = Quantity([4, 4, 4], "nanometer") @@ -38,26 +36,60 @@ def test_error_topology_mismatch(self, monkeypatch, sage_unconstrained, ethanol) topology=other_topology.to_openmm(), ) - def test_found_virtual_sites(self, monkeypatch, tip4p, water): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") + def test_found_out_of_plane_virtual_site(self, water_dimer): + pytest.importorskip("openmm") - topology = water.to_topology() - topology.box_vectors = Quantity([4, 4, 4], "nanometer") + import openmm.app + + modeller = openmm.app.Modeller( + topology=water_dimer.to_openmm(), + positions=water_dimer.get_positions().to("nanometer").to_openmm(), + ) - system = tip4p.create_openmm_system(topology) + forcefield = openmm.app.ForceField("tip5p.xml") + + modeller.addExtraParticles(forcefield=forcefield) + + system = forcefield.createSystem( + modeller.topology, + nonbondedMethod=openmm.app.PME, + nonbondedCutoff=1.0 * openmm.unit.nanometers, + constraints=openmm.app.HBonds, + rigidWater=True, + ewaldErrorTolerance=0.0005, + ) with pytest.raises( UnsupportedImportError, - match="A particle is a virtual site, which is not yet supported.", + match="A particle is a virtual site of type.*OutOfPlane.*which is not yet supported.", ): from_openmm( system=system, - topology=topology.to_openmm(), + topology=modeller.topology, ) - def test_missing_positions_warning(self, monkeypatch, sage, water): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") + @pytest.mark.skip( + reason="Need to find a way to get OpenMM to actually use TwoParticleAverageSite", + ) + def test_found_two_particle_average_virtual_site( + self, + sage_with_bond_charge, + default_integrator, + ): + simulation = sage_with_bond_charge.create_interchange( + Molecule.from_smiles("CCl").to_topology(), + ).to_openmm_simulation(integrator=default_integrator) + + with pytest.raises( + UnsupportedImportError, + match="A particle is a `TwoParticleAverage` virtual site, which is not yet supported.", + ): + from_openmm( + system=simulation.system, + topology=simulation.topology, + ) + def test_missing_positions_warning(self, sage, water): topology = water.to_topology() topology.box_vectors = Quantity([4, 4, 4], "nanometer") diff --git a/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_import.py b/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_import.py index 3646d9654..2def418f0 100644 --- a/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_import.py +++ b/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_import.py @@ -1,5 +1,6 @@ import copy import random +from collections import defaultdict import numpy import pytest @@ -20,9 +21,7 @@ @skip_if_missing("openmm") class TestFromOpenMM: - def test_simple_roundtrip(self, monkeypatch, sage_unconstrained, ethanol): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - + def test_simple_roundtrip(self, sage_unconstrained, ethanol): ethanol.generate_conformers(n_conformers=1) interchange = Interchange.from_smirnoff( @@ -69,13 +68,12 @@ def simple_system(self): @pytest.mark.parametrize("as_argument", [False, True]) def test_different_ways_to_process_box_vectors( self, - monkeypatch, as_argument, simple_system, ): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - topology = Molecule.from_smiles("C").to_topology() + topology._molecule_virtual_site_map = defaultdict(list) + topology._particle_map = {index: index for index in range(topology.n_atoms)} if as_argument: box = Interchange.from_openmm( @@ -100,8 +98,6 @@ def test_topology_and_system_box_vectors_differ( simple_system, ): """Ensure that, if box vectors specified in the topology and system differ, those in the topology are used.""" - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - topology = Molecule.from_smiles("C").to_topology() topology.box_vectors = Quantity([4, 5, 6], unit.nanometer) @@ -112,9 +108,7 @@ def test_topology_and_system_box_vectors_differ( assert numpy.diag(box.m_as(unit.nanometer)) == pytest.approx([4, 5, 6]) - def test_openmm_roundtrip_metadata(self, monkeypatch, sage): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - + def test_openmm_roundtrip_metadata(self, sage): # Make an example OpenMM Topology with metadata. # Here we use OFFTK to make the OpenMM Topology, but this could just as easily come from another source ethanol = Molecule.from_smiles("CCO") @@ -155,12 +149,10 @@ def test_openmm_roundtrip_metadata(self, monkeypatch, sage): assert atom.metadata["residue_name"] == "BNZ" @pytest.mark.slow - def test_openmm_native_roundtrip_metadata(self, monkeypatch, sage): + def test_openmm_native_roundtrip_metadata(self, sage): """ Test that metadata is the same whether we load a PDB through OpenMM+Interchange vs. Topology.from_pdb. """ - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - pdb = openmm.app.PDBFile( get_data_file_path( "ALA_GLY/ALA_GLY.pdb", @@ -184,15 +176,13 @@ def test_openmm_native_roundtrip_metadata(self, monkeypatch, sage): del off_atom_metadata["match_info"] assert roundtrip_atom.metadata == off_atom_metadata - def test_electrostatics_cutoff_not_ignored(self, monkeypatch, ethanol): + def test_electrostatics_cutoff_not_ignored(self, ethanol): pytest.importorskip("openmmforcefields") import openmm.app import openmm.unit from openmmforcefields.generators import GAFFTemplateGenerator - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - topology = ethanol.to_topology() topology.box_vectors = Quantity([4, 4, 4], "nanometer") @@ -218,13 +208,12 @@ def test_electrostatics_cutoff_not_ignored(self, monkeypatch, ethanol): assert interchange["vdW"].cutoff.m_as(unit.nanometer) == pytest.approx(1.2345) @needs_gmx - def test_fill_in_rigid_water_parameters(self, water_dimer, monkeypatch): + @pytest.mark.skip(reason="needs OpenMM -> Interchange -> GROMACS virtual sites implemented") + def test_fill_in_rigid_water_parameters(self, water_dimer): import openmm.app from openff.interchange.drivers import get_gromacs_energies - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - openmm_force_field = openmm.app.ForceField("tip3p.xml") openmm_topology = water_dimer.to_openmm() @@ -252,11 +241,12 @@ def test_fill_in_rigid_water_parameters(self, water_dimer, monkeypatch): @skip_if_missing("openmm") class TestProcessTopology: - def test_with_openff_topology(self, monkeypatch, sage, basic_top): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - + def test_with_openff_topology(self, sage, basic_top): system = sage.create_openmm_system(basic_top) + basic_top._molecule_virtual_site_map = defaultdict(list) + basic_top._particle_map = {index: index for index in range(basic_top.n_atoms)} + with_openff = Interchange.from_openmm( system=system, topology=basic_top, @@ -298,7 +288,7 @@ def test_unsupported_method(self): force.setNonbondedMethod(method) with pytest.raises(UnsupportedImportError): - _convert_nonbonded_force(force) + _convert_nonbonded_force(force, dict()) def test_parse_switching_distance(self): force = openmm.NonbondedForce() @@ -311,7 +301,7 @@ def test_parse_switching_distance(self): force.setUseSwitchingFunction(True) force.setSwitchingDistance(cutoff - switch_width) - vdw, _ = _convert_nonbonded_force(force) + vdw, _ = _convert_nonbonded_force(force=force, particle_map=dict()) assert vdw.cutoff.m_as(unit.nanometer) == pytest.approx(cutoff) assert vdw.switch_width.m_as(unit.nanometer) == pytest.approx(switch_width) @@ -324,7 +314,7 @@ def test_parse_switching_distance_unused(self): force.setCutoffDistance(cutoff) - vdw, _ = _convert_nonbonded_force(force) + vdw, _ = _convert_nonbonded_force(force=force, particle_map=dict()) assert vdw.cutoff.m_as(unit.nanometer) == pytest.approx(cutoff) assert vdw.switch_width.m_as(unit.nanometer) == 0.0 @@ -332,10 +322,8 @@ def test_parse_switching_distance_unused(self): @skip_if_missing("openmm") class TestConvertConstraints: - def test_num_constraints(self, monkeypatch, sage, basic_top): + def test_num_constraints(self, sage, basic_top): """Test that the number of constraints is preserved when converting to and from OpenMM""" - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - interchange = sage.create_interchange(basic_top) converted = Interchange.from_openmm( diff --git a/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_virtual_sites.py b/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_virtual_sites.py new file mode 100644 index 000000000..70f70a5cb --- /dev/null +++ b/openff/interchange/_tests/unit_tests/interop/openmm/_import/test_virtual_sites.py @@ -0,0 +1,102 @@ +import pytest +from openff.toolkit import Quantity, Topology + +from openff.interchange import Interchange +from openff.interchange._tests import MoleculeWithConformer +from openff.interchange.components._packmol import solvate_topology +from openff.interchange.drivers.openmm import _get_openmm_energies, get_openmm_energies + + +class TestTIP4PVirtualSites: + def test_tip4p_openmm_xml(self, water_dimer): + """ + Prepare a TIP4P water dimer with OpenMM's style of 4-site water. + + Below is used as a guide + https://openmm.github.io/openmm-cookbook/latest/notebooks/tutorials/Histone_methyltransferase_simulation_with_a_multisite_water_model_TIP4P-Ew.html + """ + pytest.importorskip("openmm") + + import openmm.app + + modeller = openmm.app.Modeller( + topology=water_dimer.to_openmm(), + positions=water_dimer.get_positions().to("nanometer").to_openmm(), + ) + + forcefield = openmm.app.ForceField("tip4pew.xml") + + modeller.addExtraParticles(forcefield=forcefield) + + system = forcefield.createSystem( + modeller.topology, + nonbondedMethod=openmm.app.PME, + nonbondedCutoff=1.0 * openmm.unit.nanometers, + constraints=openmm.app.HBonds, + rigidWater=True, + ewaldErrorTolerance=0.0005, + ) + + imported = Interchange.from_openmm( + topology=modeller.topology, + system=system, + positions=modeller.getPositions(), + box_vectors=modeller.getTopology().getPeriodicBoxVectors(), + ) + + get_openmm_energies(imported) + + @pytest.mark.skip( + reason="Rewrite to use OpenMM or update from_openmm to support `LocalCoordinatesSite`s", + ) + def test_dimer_energy_equals(self, tip4p, water_dimer): + out: Interchange = tip4p.create_interchange(water_dimer) + + roundtripped = Interchange.from_openmm( + system=out.to_openmm_system(), + topology=out.to_openmm_topology(collate=False), + positions=out.get_positions(include_virtual_sites=True).to_openmm(), + box_vectors=out.box.to_openmm(), + ) + + assert get_openmm_energies(out) == _get_openmm_energies(roundtripped) + + @pytest.mark.skip( + reason="Rewrite to use OpenMM or update from_openmm to support `LocalCoordinatesSite`s", + ) + def test_minimize_solvated_ligand(self, sage_with_tip4p, default_integrator): + topology = solvate_topology( + topology=MoleculeWithConformer.from_smiles("CO").to_topology(), + nacl_conc=Quantity(1.0, "mole / liter"), + target_density=Quantity(0.6, "gram / milliliter"), + working_directory=".", + ) + + simulation = sage_with_tip4p.create_interchange(topology).to_openmm_simulation(integrator=default_integrator) + + roundtripped = Interchange.from_openmm( + system=simulation.system, + topology=simulation.topology, + positions=simulation.context.getState(getPositions=True).getPositions(), + box_vectors=simulation.system.getDefaultPeriodicBoxVectors(), + ) + + original_energy = get_openmm_energies(simulation) + + # TODO: Much more validation could be done here, but if a simulation + # can start and minimize at all, that should catch most problems + roundtripped.minimize() + + assert get_openmm_energies(roundtripped) < original_energy + + def test_error_index_mismatch(self, tip4p, water): + out: Interchange = tip4p.create_interchange(Topology.from_molecules([water, water])) + + with pytest.raises( + ValueError, # TODO: Make a different error + match="Particle ordering mismatch", + ): + Interchange.from_openmm( + system=out.to_openmm_system(), + topology=out.to_openmm_topology(collate=True), + ) diff --git a/openff/interchange/_tests/unit_tests/operations/test_combine.py b/openff/interchange/_tests/unit_tests/operations/test_combine.py index 98cf27a00..def233b44 100644 --- a/openff/interchange/_tests/unit_tests/operations/test_combine.py +++ b/openff/interchange/_tests/unit_tests/operations/test_combine.py @@ -15,10 +15,8 @@ class TestCombine: @skip_if_missing("openmm") - def test_basic_combination(self, monkeypatch, sage_unconstrained): + def test_basic_combination(self, sage_unconstrained): """Test basic use of Interchange.__add__() based on the README example""" - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - mol = Molecule.from_smiles("C") mol.generate_conformers(n_conformers=1) top = Topology.from_molecules([mol]) @@ -40,9 +38,7 @@ def test_basic_combination(self, monkeypatch, sage_unconstrained): @pytest.mark.filterwarnings("ignore:Setting positions to None") @pytest.mark.slow - def test_parameters_do_not_clash(self, monkeypatch, sage_unconstrained): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - + def test_parameters_do_not_clash(self, sage_unconstrained): thf = Molecule.from_smiles("C1CCOC1") ace = Molecule.from_smiles("CC(=O)O") @@ -71,12 +67,10 @@ def make_interchange(molecule: Molecule) -> Interchange: # TODO: Ensure the de-duplication is maintained after exports - def test_positions_setting(self, monkeypatch, sage): + def test_positions_setting(self, sage): """Test that positions exist on the result if and only if both input objects have positions.""" - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - ethane = Molecule.from_smiles("CC") methane = Molecule.from_smiles("C") @@ -98,13 +92,10 @@ def test_positions_setting(self, monkeypatch, sage): @pytest.mark.parametrize("handler", ["Electrostatics", "vdW"]) def test_error_mismatched_cutoffs( self, - monkeypatch, sage, basic_top, handler, ): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - sage_modified = ForceField("openff-2.1.0.offxml") sage_modified[handler].cutoff *= 1.5 @@ -119,12 +110,9 @@ def test_error_mismatched_cutoffs( def test_error_mismatched_switching_function( self, - monkeypatch, sage, basic_top, ): - monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1") - sage_modified = ForceField("openff-2.1.0.offxml") sage_modified["vdW"].switch_width *= 0.0 diff --git a/openff/interchange/_tests/unit_tests/test_models.py b/openff/interchange/_tests/unit_tests/test_models.py index 6c3819925..caecfeb4d 100644 --- a/openff/interchange/_tests/unit_tests/test_models.py +++ b/openff/interchange/_tests/unit_tests/test_models.py @@ -1,5 +1,8 @@ +import re + from openff.interchange.models import ( AngleKey, + BaseVirtualSiteKey, BondKey, ImproperTorsionKey, PotentialKey, @@ -197,3 +200,9 @@ def test_torsionkey_eq_hash(): None, 1.5, ) + + +def test_virtual_site_key_repr(): + key = BaseVirtualSiteKey(orientation_atom_indices=[3, 4, 5], type="fOO", name="bAR") + + assert re.search("orientation atom ind.*3, 4, 5", repr(key)) diff --git a/openff/interchange/components/_packmol.py b/openff/interchange/components/_packmol.py index c9de1ace5..318eb6b62 100644 --- a/openff/interchange/components/_packmol.py +++ b/openff/interchange/components/_packmol.py @@ -567,7 +567,7 @@ def pack_box( box_shape: ArrayLike = RHOMBIC_DODECAHEDRON, center_solute: bool | Literal["BOX_VECS", "ORIGIN", "BRICK"] = False, working_directory: str | None = None, - retain_working_files: bool = False, + retain_working_files: bool = True, ) -> Topology: """ Run packmol to generate a box containing a mixture of molecules. @@ -816,6 +816,7 @@ def solvate_topology( box_shape: NDArray = RHOMBIC_DODECAHEDRON, target_density: Quantity = Quantity(0.9, "gram / milliliter"), tolerance: Quantity = Quantity(2.0, "angstrom"), + working_directory: str | None = None, ) -> Topology: """ Add water and ions to neutralise and solvate a topology. @@ -908,6 +909,7 @@ def solvate_topology( solute=topology, tolerance=tolerance, box_vectors=box_vectors, + working_directory=working_directory, ) @@ -918,6 +920,7 @@ def solvate_topology_nonwater( padding: Quantity = Quantity(1.2, "nanometer"), box_shape: NDArray = RHOMBIC_DODECAHEDRON, tolerance: Quantity = Quantity(2.0, "angstrom"), + working_directory: str | None = None, ) -> Topology: """ Add water and ions to neutralise and solvate a topology. @@ -989,4 +992,5 @@ def solvate_topology_nonwater( tolerance=tolerance, box_vectors=box_vectors, center_solute=True, + working_directory=working_directory, ) diff --git a/openff/interchange/components/_particles.py b/openff/interchange/components/_particles.py index aff3b4fc1..caf487a21 100644 --- a/openff/interchange/components/_particles.py +++ b/openff/interchange/components/_particles.py @@ -10,6 +10,8 @@ from openff.interchange.pydantic import _BaseModel +# TODO: This is actually quite SMIRNOFF-specific, should be refactored out since +# it can't really represent a generic/abstract base class for virtual sites class _VirtualSite(_BaseModel, abc.ABC): type: str distance: _DistanceQuantity diff --git a/openff/interchange/components/interchange.py b/openff/interchange/components/interchange.py index ef2b71de7..7380bb071 100644 --- a/openff/interchange/components/interchange.py +++ b/openff/interchange/components/interchange.py @@ -904,6 +904,8 @@ def from_openmm( representing an arbitrarily harmonic angle. It is expected that these values will be overwritten by runtime MD options. + See more limitations and sharp edges in the user guide: https://docs.openforcefield.org/projects/interchange/en/latest/using/edges.html#quirks-of-from-openmm + Parameters ---------- system : openmm.System diff --git a/openff/interchange/components/toolkit.py b/openff/interchange/components/toolkit.py index 361597550..121beef60 100644 --- a/openff/interchange/components/toolkit.py +++ b/openff/interchange/components/toolkit.py @@ -1,6 +1,7 @@ """Utilities for processing and interfacing with the OpenFF Toolkit.""" -from typing import TYPE_CHECKING, Union +from collections import defaultdict +from typing import TYPE_CHECKING import networkx import numpy @@ -10,6 +11,8 @@ from openff.toolkit.utils.collections import ValidatedList from openff.utilities.utilities import has_package +from openff.interchange.models import ImportedVirtualSiteKey + if has_package("openmm") or TYPE_CHECKING: import openmm.app @@ -25,7 +28,7 @@ def _get_num_h_bonds(topology: "Topology") -> int: return n_bonds_containing_hydrogen -def _get_14_pairs(topology_or_molecule: Union["Topology", "Molecule"]): +def _get_14_pairs(topology_or_molecule: Topology | Molecule): """Generate tuples of atom pairs, including symmetric duplicates.""" # TODO: A replacement of Topology.nth_degree_neighbors in the toolkit # may implement this in the future. @@ -105,27 +108,74 @@ def _check_electrostatics_handlers(force_field: "ForceField") -> bool: return False -def _simple_topology_from_openmm(openmm_topology: "openmm.app.Topology") -> Topology: - """Convert an OpenMM Topology into an OpenFF Topology consisting **only** of so-called `_SimpleMolecule`s.""" +def _simple_topology_from_openmm( + openmm_topology: "openmm.app.Topology", + system: openmm.System | None = None, +) -> Topology: + """ + Convert an OpenMM Topology into an OpenFF Topology consisting **only** of so-called `_SimpleMolecule`s. + + Arguments + --------- + openmm_topology: openmm.app.Topology + The OpenMM Topology to convert. + system: openmm.System, optional + The OpenMM System associated with the topology. Only needed if there are virtual sites in the topology. + """ # TODO: Splice in fully-defined OpenFF `Molecule`s? graph = networkx.Graph() - # TODO: This is nearly identical to Topology._openmm_topology_to_networkx. - # Should this method be replaced with a direct call to that? + virtual_sites: list[ImportedVirtualSiteKey] = list() + + # map indices of OpenMM system with virtual site to topology indices, or virtual site keys, + # in associated OpenFF topology, since stripping out virtual sites offsets particle indices + openmm_openff_particle_map: dict[int, int | ImportedVirtualSiteKey] = dict() + for atom in openmm_topology.atoms(): - graph.add_node( - atom.index, - atomic_number=atom.element.atomic_number, - name=atom.name, - residue_name=atom.residue.name, - # Note that residue number is mapped to residue.id here. The use of id vs. number varies in other packages - # and the convention for the OpenFF-OpenMM interconversion is recorded at - # https://docs.openforcefield.org/projects/toolkit/en/0.15.1/users/molecule_conversion.html - residue_number=atom.residue.id, - insertion_code=atom.residue.insertionCode, - chain_id=atom.residue.chain.id, - ) + if atom.element is None: + assert isinstance( + system, + openmm.System, + ), "`system` argument required if virtual sites are present in the topology" + try: + # assume ThreeParticleAverageSite for now + orientation_atom_indices = [ + openmm_openff_particle_map[system.getVirtualSite(atom.index).getParticle(i)] for i in range(3) + ] + except openmm.OpenMMException as error: + if "This particle is not a virtual site" in str(error): + raise ValueError( + "Particle ordering mismatch between OpenMM system and topology. " + f"Look at particle {atom.index} in the topology.", + ) from error + virtual_sites.append( + ImportedVirtualSiteKey( + orientation_atom_indices=orientation_atom_indices, + name=atom.name, + type="ThreeParticleAverageSite", + ), + ) + + # TODO: This will break if virtual sites are not after their parent/orientation atoms + openmm_openff_particle_map[atom.index] = virtual_sites[-1] + + else: + graph.add_node( + atom.index, + atomic_number=getattr(atom.element, "atomic_number", 0), + name=atom.name, + residue_name=atom.residue.name, + # Note that residue number is mapped to residue.id here. The use of id vs. number + # varies in other packages and the convention for the OpenFF-OpenMM interconversion + # is recorded at + # https://docs.openforcefield.org/projects/toolkit/en/0.15.1/users/molecule_conversion.html + residue_number=atom.residue.id, + insertion_code=atom.residue.insertionCode, + chain_id=atom.residue.chain.id, + ) + + openmm_openff_particle_map[atom.index] = atom.index - len(virtual_sites) for bond in openmm_topology.bonds(): graph.add_edge( @@ -133,7 +183,21 @@ def _simple_topology_from_openmm(openmm_topology: "openmm.app.Topology") -> Topo bond.atom2.index, ) - return _simple_topology_from_graph(graph) + topology = _simple_topology_from_graph(graph) + + topology._molecule_virtual_site_map = defaultdict(list) + + # TODO: This iteration strategy scales horribly with system size - need to refactor - + # since looking up topology atom indices is slow. It's probably repetitive to + # look up the molecule index over and over again + for particle in virtual_sites: + molecule_index = topology.molecule_index(topology.atom(particle.orientation_atom_indices[0]).molecule) + + topology._molecule_virtual_site_map[molecule_index].append(particle) + + topology._particle_map = openmm_openff_particle_map + + return topology def _simple_topology_from_graph(graph: networkx.Graph) -> Topology: @@ -147,7 +211,10 @@ def _simple_topology_from_graph(graph: networkx.Graph) -> Topology: # the subgraphs are returned out of "atom order", like # if atoms in an later molecule have lesser atom indices # than this molecule - assert topology.n_atoms == next(iter(subgraph.nodes)) + # + # Oct 2024 - need to add a test case for above? + # these values are not necessarily equal because of virtual sites + assert topology.n_atoms <= next(iter(subgraph.nodes)) topology.add_molecule(_SimpleMolecule._from_subgraph(subgraph)) diff --git a/openff/interchange/interop/_virtual_sites.py b/openff/interchange/interop/_virtual_sites.py index 674d0f560..296355913 100644 --- a/openff/interchange/interop/_virtual_sites.py +++ b/openff/interchange/interop/_virtual_sites.py @@ -2,22 +2,69 @@ Common helpers for exporting virtual sites. """ +import abc from collections import defaultdict +from typing import TYPE_CHECKING import numpy from openff.toolkit import Quantity, unit -from openff.interchange import Interchange from openff.interchange.exceptions import ( MissingPositionsError, MissingVirtualSitesError, ) -from openff.interchange.models import VirtualSiteKey +from openff.interchange.models import BaseVirtualSiteKey +from openff.interchange.pydantic import _BaseModel + +if TYPE_CHECKING: + import openmm + + from openff.interchange import Interchange + + +class _OpenMMVirtualSite(_BaseModel, abc.ABC): + particles: list[int] + + @abc.abstractmethod + def to_openmm(self) -> "openmm.VirtualSite": + """Create a (subclass of) openmm.VirtualSite""" + raise NotImplementedError() + + +class _TwoParticleAverageSite(_OpenMMVirtualSite): + weights: list[float] + + def to_openmm(self): + import openmm + + return openmm.TwoParticleAverageSite( + particle1=self.particles[0], + particle2=self.particles[1], + weight1=self.weights[0], + weight2=self.weights[1], + ) + + +class _ThreeParticleAverageSite(_OpenMMVirtualSite): + weights: list[float] + + def to_openmm(self): + import openmm + + # These arguments can't be named because (vague SWIG reasons) + return openmm.ThreeParticleAverageSite( + self.particles[0], # particle1 + self.particles[1], # particle2 + self.particles[2], # particle3 + self.weights[0], # weight1 + self.weights[1], # weight2 + self.weights[2], # weight3 + ) def _virtual_site_parent_molecule_mapping( - interchange: Interchange, -) -> dict[VirtualSiteKey, int]: + interchange: "Interchange", +) -> dict[BaseVirtualSiteKey, int]: """ Map `VirtualSiteKey`s the index of the molecule they belong to. @@ -32,7 +79,7 @@ def _virtual_site_parent_molecule_mapping( A dictionary mapping virtual site keys to the index of the molecule they belong to. """ - mapping: dict[VirtualSiteKey, int] = dict() + mapping: dict[BaseVirtualSiteKey, int] = dict() if "VirtualSites" not in interchange.collections: return mapping @@ -41,7 +88,7 @@ def _virtual_site_parent_molecule_mapping( # how they are presented in the iterator; this may cause problems when a # molecule (a large ligand? polymer?) has many virtual sites for virtual_site_key in interchange["VirtualSites"].key_map: - assert isinstance(virtual_site_key, VirtualSiteKey) + assert isinstance(virtual_site_key, BaseVirtualSiteKey), type(virtual_site_key) parent_atom_index = virtual_site_key.orientation_atom_indices[0] @@ -54,8 +101,19 @@ def _virtual_site_parent_molecule_mapping( return mapping +def _get_molecule_virtual_site_map(interchange: "Interchange") -> defaultdict[int, list[BaseVirtualSiteKey]]: + virtual_site_molecule_map = _virtual_site_parent_molecule_mapping(interchange) + + molecule_virtual_site_map = defaultdict(list) + + for virtual_site, molecule_index in virtual_site_molecule_map.items(): + molecule_virtual_site_map[molecule_index].append(virtual_site) + + return molecule_virtual_site_map + + def get_positions_with_virtual_sites( - interchange: Interchange, + interchange: "Interchange", collate: bool = False, use_zeros: bool = False, ) -> Quantity: @@ -74,7 +132,7 @@ def get_positions_with_virtual_sites( raise MissingVirtualSitesError() # map of molecule index to *list* of virtual site keys contained therein - molecule_virtual_site_map: defaultdict[int, list[VirtualSiteKey]] = defaultdict( + molecule_virtual_site_map: defaultdict[int, list[BaseVirtualSiteKey]] = defaultdict( list, ) @@ -161,7 +219,7 @@ def get_positions_with_virtual_sites( def _get_separation_by_atom_indices( - interchange: Interchange, + interchange: "Interchange", atom_indices: tuple[int, ...], prioritize_geometry: bool = False, ) -> Quantity: diff --git a/openff/interchange/interop/common.py b/openff/interchange/interop/common.py index 140245078..3a5866ed7 100644 --- a/openff/interchange/interop/common.py +++ b/openff/interchange/interop/common.py @@ -8,7 +8,7 @@ MissingVirtualSitesError, UnsupportedExportError, ) -from openff.interchange.models import VirtualSiteKey +from openff.interchange.models import BaseVirtualSiteKey from openff.interchange.smirnoff import SMIRNOFFVirtualSiteCollection @@ -74,14 +74,14 @@ def _build_particle_map( interchange: Interchange, molecule_virtual_site_map, collate: bool = False, -) -> dict[int | VirtualSiteKey, int]: +) -> dict[int | BaseVirtualSiteKey, int]: """ Build a dict mapping particle indices between a topology and another object. If `collate=True`, virtual sites are collated with each molecule's atoms. If `collate=False`, virtual sites go at the very end, after all atoms were added. """ - particle_map: dict[int | VirtualSiteKey, int] = dict() + particle_map: dict[int | BaseVirtualSiteKey, int] = dict() particle_index = 0 diff --git a/openff/interchange/interop/gromacs/export/_virtual_sites.py b/openff/interchange/interop/gromacs/export/_virtual_sites.py index ccfe2e672..eddce3782 100644 --- a/openff/interchange/interop/gromacs/export/_virtual_sites.py +++ b/openff/interchange/interop/gromacs/export/_virtual_sites.py @@ -13,7 +13,7 @@ GROMACSVirtualSite3, GROMACSVirtualSite3fad, ) -from openff.interchange.models import VirtualSiteKey +from openff.interchange.models import BaseVirtualSiteKey from openff.interchange.smirnoff._virtual_sites import ( _BondChargeVirtualSite, _DivalentLonePairVirtualSite, @@ -25,8 +25,8 @@ def _create_gromacs_virtual_site( interchange: Interchange, virtual_site: "_VirtualSite", - virtual_site_key: VirtualSiteKey, - particle_map: dict[int | VirtualSiteKey, int], + virtual_site_key: BaseVirtualSiteKey, + particle_map: dict[int | BaseVirtualSiteKey, int], ) -> GROMACSVirtualSite: # Orientation atom indices are topology indices, but here they need to be indexed as molecule # indices. Store the difference between an orientation atom's molecule and topology indices. diff --git a/openff/interchange/interop/openmm/_import/_import.py b/openff/interchange/interop/openmm/_import/_import.py index fd4fd7686..5c8cf28a7 100644 --- a/openff/interchange/interop/openmm/_import/_import.py +++ b/openff/interchange/interop/openmm/_import/_import.py @@ -2,7 +2,10 @@ from typing import TYPE_CHECKING, Union from openff.toolkit import Quantity, Topology +from openff.units.openmm import ensure_quantity +from openff.units.openmm import from_openmm as from_openmm_ from openff.utilities.utilities import has_package, requires_package +from pydantic import ValidationError from openff.interchange.common._nonbonded import vdWCollection from openff.interchange.common._valence import ( @@ -15,7 +18,9 @@ from openff.interchange.interop.openmm._import._nonbonded import ( BasicElectrostaticsCollection, ) +from openff.interchange.interop.openmm._import._virtual_sites import _convert_virtual_sites from openff.interchange.interop.openmm._import.compat import _check_compatible_inputs +from openff.interchange.models import ImportedVirtualSiteKey from openff.interchange.warnings import MissingPositionsWarning if has_package("openmm"): @@ -45,11 +50,9 @@ def from_openmm( _check_compatible_inputs(system=system, topology=topology) if isinstance(topology, openmm.app.Topology): - from openff.units.openmm import from_openmm as from_openmm_ - from openff.interchange.components.toolkit import _simple_topology_from_openmm - openff_topology = _simple_topology_from_openmm(topology) + openff_topology = _simple_topology_from_openmm(topology, system) if topology.getPeriodicBoxVectors() is not None: openff_topology.box_vectors = from_openmm_(topology.getPeriodicBoxVectors()) @@ -70,32 +73,41 @@ def from_openmm( interchange = Interchange(topology=openff_topology) - if system: - constraints = _convert_constraints(system) - - if constraints is not None: - interchange.collections["Constraints"] = constraints - - for force in system.getForces(): - if isinstance(force, openmm.NonbondedForce): - vdw, coul = _convert_nonbonded_force(force) - interchange.collections["vdW"] = vdw - interchange.collections["Electrostatics"] = coul - elif isinstance(force, openmm.HarmonicBondForce): - bonds = _convert_harmonic_bond_force(force) - interchange.collections["Bonds"] = bonds - elif isinstance(force, openmm.HarmonicAngleForce): - angles = _convert_harmonic_angle_force(force) - interchange.collections["Angles"] = angles - elif isinstance(force, openmm.PeriodicTorsionForce): - proper_torsions = _convert_periodic_torsion_force(force) - interchange.collections["ProperTorsions"] = proper_torsions - elif isinstance(force, openmm.CMMotionRemover): - pass - else: - raise UnsupportedImportError( - f"Unsupported OpenMM Force type ({type(force)}) found.", - ) + interchange.topology._molecule_virtual_site_map = openff_topology._molecule_virtual_site_map + interchange.topology._particle_map = openff_topology._particle_map + + # TODO: Actually build up the VirtualSiteCollection, maybe using _molecule_virtual_site_map + + constraints = _convert_constraints(system, openff_topology._particle_map) + + if constraints is not None: + interchange.collections["Constraints"] = constraints + + virtual_sites = _convert_virtual_sites(system, openff_topology) + + if virtual_sites is not None: + interchange.collections["VirtualSites"] = virtual_sites + + for force in system.getForces(): + if isinstance(force, openmm.NonbondedForce): + vdw, coul = _convert_nonbonded_force(force, openff_topology._particle_map) + interchange.collections["vdW"] = vdw + interchange.collections["Electrostatics"] = coul + elif isinstance(force, openmm.HarmonicBondForce): + bonds = _convert_harmonic_bond_force(force) + interchange.collections["Bonds"] = bonds + elif isinstance(force, openmm.HarmonicAngleForce): + angles = _convert_harmonic_angle_force(force) + interchange.collections["Angles"] = angles + elif isinstance(force, openmm.PeriodicTorsionForce): + proper_torsions = _convert_periodic_torsion_force(force) + interchange.collections["ProperTorsions"] = proper_torsions + elif isinstance(force, openmm.CMMotionRemover): + pass + else: + raise UnsupportedImportError( + f"Unsupported OpenMM Force type ({type(force)}) found.", + ) if positions is None: warnings.warn( @@ -104,7 +116,11 @@ def from_openmm( ) else: - interchange.positions = positions + assert len(positions) == len(interchange.topology._particle_map) + + interchange.positions = ensure_quantity(positions, "openff")[ + [key for key, val in interchange.topology._particle_map.items() if isinstance(val, int)] + ] if box_vectors is not None: _box_vectors = box_vectors @@ -115,8 +131,6 @@ def from_openmm( else: # If there is no box argument passed and the topology is non-periodic # and the system does not have default box vectors, it'll end up as None - from openff.units.openmm import from_openmm as from_openmm_ - _box_vectors = from_openmm_(system.getDefaultPeriodicBoxVectors()) # TODO: Does this run through the Interchange.box validator? @@ -132,6 +146,7 @@ def from_openmm( def _convert_constraints( system: "openmm.System", + particle_map: dict[int, int | ImportedVirtualSiteKey], ) -> ConstraintCollection | None: from openff.toolkit import unit @@ -170,13 +185,14 @@ def _convert_constraints( distance = _distance.value_in_unit(openmm.unit.nanometer) - constraints.key_map[BondKey(atom_indices=(atom1, atom2))] = _keys[distance] + constraints.key_map[BondKey(atom_indices=(particle_map[atom1], particle_map[atom2]))] = _keys[distance] return constraints def _convert_nonbonded_force( force: "openmm.NonbondedForce", + particle_map: dict[int, int | ImportedVirtualSiteKey], ) -> tuple[vdWCollection, BasicElectrostaticsCollection]: from openff.units.openmm import from_openmm as from_openmm_quantity @@ -196,7 +212,10 @@ def _convert_nonbonded_force( for idx in range(n_parametrized_particles): charge, sigma, epsilon = force.getParticleParameters(idx) - top_key = TopologyKey(atom_indices=(idx,)) + try: + top_key = TopologyKey(atom_indices=(particle_map[idx],)) + except ValidationError: + top_key: ImportedVirtualSiteKey = particle_map[idx] # type: ignore[no-redef] pot = Potential( parameters={ diff --git a/openff/interchange/interop/openmm/_import/_nonbonded.py b/openff/interchange/interop/openmm/_import/_nonbonded.py index 53c3a6fac..3a3c6b9e8 100644 --- a/openff/interchange/interop/openmm/_import/_nonbonded.py +++ b/openff/interchange/interop/openmm/_import/_nonbonded.py @@ -33,9 +33,6 @@ def _get_charges( # type: ignore for topology_key, potential_key in self.key_map.items(): potential = self.potentials[potential_key] - if len(topology_key.atom_indices) * len(potential.parameters) != 1: - raise RuntimeError() - charges[topology_key] = potential.parameters["charge"] return charges diff --git a/openff/interchange/interop/openmm/_import/_virtual_sites.py b/openff/interchange/interop/openmm/_import/_virtual_sites.py new file mode 100644 index 000000000..b90317680 --- /dev/null +++ b/openff/interchange/interop/openmm/_import/_virtual_sites.py @@ -0,0 +1,83 @@ +from typing import Literal + +import openmm +from openff.toolkit import Quantity, Topology +from pydantic import Field + +from openff.interchange.components.potentials import Collection, Potential +from openff.interchange.models import PotentialKey, VirtualSiteKey + + +class _VirtualSiteCollection(Collection): + """ + A handler which stores the information necessary to construct non-SMIRNOFF virtual sites. + """ + + key_map: dict[VirtualSiteKey, PotentialKey] = Field( + dict(), + description="A mapping between VirtualSiteKey objects and PotentialKey objects.", + ) # type: ignore[assignment] + + type: Literal["VirtualSites"] = "VirtualSites" + expression: Literal[""] = "" + virtual_site_key_topology_index_map: dict[VirtualSiteKey, int] = Field( + dict(), + description="A mapping between VirtualSiteKey objects (stored analogously to TopologyKey objects" + "in other handlers) and topology indices describing the associated virtual site", + ) + + exclusion_policy: Literal["parents"] = "parents" + + +def _convert_virtual_sites( + system: openmm.System, + topology: Topology, +) -> _VirtualSiteCollection | None: + if topology._particle_map is None: + return None + + collection = _VirtualSiteCollection() + + for particle_index in range(system.getNumParticles()): + if system.getParticleMass(particle_index)._value != 0.0: + continue + + virtual_site = system.getVirtualSite(particle_index) + + openmm_particle_indices: tuple[int, int, int] = ( + virtual_site.getParticle(0), + virtual_site.getParticle(1), + virtual_site.getParticle(2), + ) + + # get ImportedVirtualSiteKey from topology._molecule_virtual_site_map + for _, list_of_virtual_sites in topology._molecule_virtual_site_map.items(): + for virtual_site_key in list_of_virtual_sites: + if virtual_site_key.orientation_atom_indices == tuple( + topology._particle_map[i] for i in openmm_particle_indices + ): + # TODO: Not the cleanest way to de-duplicate virtual site parameters, + # but other info (i.e. atom name) is insufficient + id = type(virtual_site).__name__ + "-".join( + map( + str, + [ + virtual_site.getWeight(0), + virtual_site.getWeight(1), + virtual_site.getWeight(2), + ], + ), + ) + + potential_key = PotentialKey(id=id, virtual_site_type="ThreeParticleSite") + + collection.key_map[virtual_site_key] = potential_key + # TODO: This is where geometry-specific information is stored in SMIRNOFF virtual sites, + # but in this case that's not needed + collection.potentials[potential_key] = Potential( + parameters={ + "weights": Quantity([virtual_site.getWeight(i) for i in range(3)], "dimensionless"), + }, + ) + + return collection diff --git a/openff/interchange/interop/openmm/_import/compat.py b/openff/interchange/interop/openmm/_import/compat.py index fc54cdebe..c8b091c48 100644 --- a/openff/interchange/interop/openmm/_import/compat.py +++ b/openff/interchange/interop/openmm/_import/compat.py @@ -4,6 +4,11 @@ from openff.interchange.exceptions import UnsupportedImportError +_SUPPORTED_VIRTUAL_SITE_TYPES: set[type] = { + openmm.ThreeParticleAverageSite, + openmm.LocalCoordinatesSite, +} + def _check_compatible_inputs( system: openmm.System, @@ -12,9 +17,11 @@ def _check_compatible_inputs( """Check that inputs are compatible and supported.""" for index in range(system.getNumParticles()): if system.isVirtualSite(index): - raise UnsupportedImportError( - "A particle is a virtual site, which is not yet supported.", - ) + if type(system.getVirtualSite(index)) not in _SUPPORTED_VIRTUAL_SITE_TYPES: + raise UnsupportedImportError( + f"A particle is a virtual site of type {type(system.getVirtualSite(index))}, " + "which is not yet supported.", + ) if isinstance(topology, Topology): _topology = topology.to_openmm() diff --git a/openff/interchange/interop/openmm/_nonbonded.py b/openff/interchange/interop/openmm/_nonbonded.py index a4817022b..2dffb21b0 100644 --- a/openff/interchange/interop/openmm/_nonbonded.py +++ b/openff/interchange/interop/openmm/_nonbonded.py @@ -22,9 +22,10 @@ ) from openff.interchange.interop.common import _build_particle_map from openff.interchange.models import ( + BaseVirtualSiteKey, + ImportedVirtualSiteKey, SingleAtomChargeTopologyKey, TopologyKey, - VirtualSiteKey, ) from openff.interchange.warnings import NonbondedSettingsWarning @@ -56,7 +57,7 @@ def _process_nonbonded_forces( system: openmm.System, combine_nonbonded_forces: bool = False, ewald_tolerance: float = 1e-4, -) -> dict[int | VirtualSiteKey, int]: +) -> dict[int | BaseVirtualSiteKey, int]: """ Process the non-bonded collections in an Interchange into corresponding openmm objects. @@ -91,11 +92,11 @@ def _process_nonbonded_forces( _check_virtual_site_exclusion_policy(virtual_sites) virtual_site_molecule_map: dict[ - VirtualSiteKey, + BaseVirtualSiteKey, int, ] = _virtual_site_parent_molecule_mapping(interchange) - molecule_virtual_site_map: dict[int, list[VirtualSiteKey]] = defaultdict(list) + molecule_virtual_site_map: dict[int, list[BaseVirtualSiteKey]] = defaultdict(list) for virtual_site_key, molecule_index in virtual_site_molecule_map.items(): molecule_virtual_site_map[molecule_index].append(virtual_site_key) @@ -165,7 +166,7 @@ def _add_particles_to_system( interchange: "Interchange", system: openmm.System, molecule_virtual_site_map, -) -> dict[int | VirtualSiteKey, int]: +) -> dict[int | BaseVirtualSiteKey, int]: particle_map = _build_particle_map( interchange, molecule_virtual_site_map, @@ -281,8 +282,8 @@ def _create_single_nonbonded_force( interchange: "Interchange", system: openmm.System, ewald_tolerance: float, - molecule_virtual_site_map: dict["Molecule", list[VirtualSiteKey]], - openff_openmm_particle_map: dict[int | VirtualSiteKey, int], + molecule_virtual_site_map: dict["Molecule", list[BaseVirtualSiteKey]], + openff_openmm_particle_map: dict[int | BaseVirtualSiteKey, int], ): """Create a single openmm.NonbondedForce from vdW/electrostatics/virtual site collections.""" if data.mixing_rule not in ("lorentz-berthelot", ""): @@ -449,8 +450,17 @@ def _create_single_nonbonded_force( "vdW and/or electrostatics interactions", ) - charge_increments = coul.potentials[coul_key].parameters["charge_increments"] - charge = to_openmm_quantity(-sum(charge_increments)) + try: + # SMIRNOFF-style virtual sites are charged from charge increments + # moved from their parent atoms + charge_increments = coul.potentials[coul_key].parameters["charge_increments"] + + charge = to_openmm_quantity(-sum(charge_increments)) + except KeyError: + # but stuff imported from other sources just has the charge right there + assert isinstance(virtual_site_key, ImportedVirtualSiteKey) + + charge = coul.potentials[coul_key].parameters["charge"].to_openmm() vdw_parameters = vdw.potentials[vdw_key].parameters sigma = to_openmm_quantity(vdw_parameters["sigma"]) @@ -464,7 +474,13 @@ def _create_single_nonbonded_force( f"Mismatch in system ({system_index=}) and force ({force_index=}) indexing", ) - parent_atom_index = openff_openmm_particle_map[virtual_site_object.orientations[0]] + try: + parent_atom_index = openff_openmm_particle_map[virtual_site_object.orientations[0]] + except AttributeError: + # TODO: Should unify these objects + assert isinstance(virtual_site_key, ImportedVirtualSiteKey) + + parent_atom_index = openff_openmm_particle_map[virtual_site_object.particles[0]] # type: ignore[attr-defined] parent_virtual_particle_mapping[parent_atom_index].append(force_index) @@ -572,7 +588,7 @@ def _create_multiple_nonbonded_forces( system: openmm.System, ewald_tolerance: float, molecule_virtual_site_map: dict, - openff_openmm_particle_map: dict[int | VirtualSiteKey, int], + openff_openmm_particle_map: dict[int | BaseVirtualSiteKey, int], ): from openff.interchange.components.toolkit import _get_14_pairs @@ -745,7 +761,7 @@ def _create_multiple_nonbonded_forces( def _create_vdw_force( data: _NonbondedData, interchange: "Interchange", - molecule_virtual_site_map: dict[int, list[VirtualSiteKey]], + molecule_virtual_site_map: dict[int, list[BaseVirtualSiteKey]], has_virtual_sites: bool, ) -> openmm.CustomNonbondedForce | None: vdw_collection: vdWCollection | None = data.vdw_collection @@ -836,7 +852,7 @@ def _create_electrostatics_force( data: _NonbondedData, interchange: "Interchange", ewald_tolerance: float, - molecule_virtual_site_map: dict[int, list[VirtualSiteKey]], + molecule_virtual_site_map: dict[int, list[BaseVirtualSiteKey]], has_virtual_sites: bool, openff_openmm_particle_map, ) -> openmm.NonbondedForce | None: @@ -934,8 +950,8 @@ def _set_particle_parameters( electrostatics_force: openmm.NonbondedForce, interchange: "Interchange", has_virtual_sites: bool, - molecule_virtual_site_map: dict[int, list[VirtualSiteKey]], - openff_openmm_particle_map: dict[int | VirtualSiteKey, int], + molecule_virtual_site_map: dict[int, list[BaseVirtualSiteKey]], + openff_openmm_particle_map: dict[int | BaseVirtualSiteKey, int], ): # TODO: Some funky plugins might have no charges, so there eventually should be # handling for electrostatics_force = None diff --git a/openff/interchange/interop/openmm/_topology.py b/openff/interchange/interop/openmm/_topology.py index dd2b4f8b8..bbb132c24 100644 --- a/openff/interchange/interop/openmm/_topology.py +++ b/openff/interchange/interop/openmm/_topology.py @@ -7,7 +7,7 @@ from openff.utilities.utilities import has_package from openff.interchange import Interchange -from openff.interchange.models import VirtualSiteKey +from openff.interchange.models import BaseVirtualSiteKey if has_package("openmm") or TYPE_CHECKING: import openmm.app @@ -34,27 +34,18 @@ def to_openmm_topology( # Heavily cribbed from the toolkit # https://github.com/openforcefield/openff-toolkit/blob/0.11.0rc2/openff/toolkit/topology/topology.py - from collections import defaultdict - from openff.toolkit import Topology from openff.toolkit.topology._mm_molecule import _SimpleBond from openff.toolkit.topology.molecule import Bond - from openff.interchange.interop._virtual_sites import ( - _virtual_site_parent_molecule_mapping, - ) + from openff.interchange.interop._virtual_sites import _get_molecule_virtual_site_map # Copy topology to avoid modifying input (eg, when generating atom names) topology = Topology(interchange.topology) - virtual_site_molecule_map = _virtual_site_parent_molecule_mapping(interchange) - - molecule_virtual_site_map = defaultdict(list) + molecule_virtual_site_map = _get_molecule_virtual_site_map(interchange) - for virtual_site, molecule_index in virtual_site_molecule_map.items(): - molecule_virtual_site_map[molecule_index].append(virtual_site) - - has_virtual_sites = len(virtual_site_molecule_map) > 0 + has_virtual_sites = len(molecule_virtual_site_map) > 0 virtual_site_element = openmm.app.element.Element.getByMass(0) @@ -128,7 +119,7 @@ def to_openmm_topology( last_residue = residue if has_virtual_sites and collate: - virtual_sites_in_this_molecule: list[VirtualSiteKey] = molecule_virtual_site_map[molecule_index] + virtual_sites_in_this_molecule: list[BaseVirtualSiteKey] = molecule_virtual_site_map[molecule_index] for this_virtual_site in virtual_sites_in_this_molecule: virtual_site_name = this_virtual_site.name @@ -178,7 +169,8 @@ def to_openmm_topology( chain = None residue = None for virtual_site_key in interchange["VirtualSites"].key_map: - assert isinstance(virtual_site_key, VirtualSiteKey) + assert isinstance(virtual_site_key, BaseVirtualSiteKey) + parent_atom_index = virtual_site_key.orientation_atom_indices[0] parent_atom = omm_atoms[parent_atom_index] diff --git a/openff/interchange/interop/openmm/_valence.py b/openff/interchange/interop/openmm/_valence.py index f8eb40f63..77a4374f5 100644 --- a/openff/interchange/interop/openmm/_valence.py +++ b/openff/interchange/interop/openmm/_valence.py @@ -7,7 +7,7 @@ from openff.utilities.utilities import has_package from openff.interchange.exceptions import UnsupportedExportError -from openff.interchange.models import VirtualSiteKey +from openff.interchange.models import BaseVirtualSiteKey if has_package("openmm"): import openmm @@ -16,7 +16,7 @@ def _process_constraints( interchange, openmm_sys, - particle_map: dict[int | VirtualSiteKey, int], + particle_map: dict[int | BaseVirtualSiteKey, int], ) -> set[tuple[int, ...]]: """ Process the Constraints section of an Interchange object. @@ -30,6 +30,7 @@ def _process_constraints( for top_key, pot_key in constraint_handler.key_map.items(): openff_indices = top_key.atom_indices + openmm_indices = tuple(particle_map[index] for index in openff_indices) params = constraint_handler.potentials[pot_key].parameters @@ -51,7 +52,7 @@ def _process_bond_forces( openmm_sys, add_constrained_forces: bool, constrained_pairs: set[tuple[int, ...]], - particle_map: dict[int | VirtualSiteKey, int], + particle_map: dict[int | BaseVirtualSiteKey, int], ): """ Process the Bonds section of an Interchange object. diff --git a/openff/interchange/interop/openmm/_virtual_sites.py b/openff/interchange/interop/openmm/_virtual_sites.py index 5f36572b5..72304397b 100644 --- a/openff/interchange/interop/openmm/_virtual_sites.py +++ b/openff/interchange/interop/openmm/_virtual_sites.py @@ -6,7 +6,8 @@ from openff.interchange import Interchange from openff.interchange.components._particles import _VirtualSite -from openff.interchange.models import VirtualSiteKey +from openff.interchange.interop._virtual_sites import _OpenMMVirtualSite +from openff.interchange.models import BaseVirtualSiteKey if has_package("openmm"): import openmm @@ -14,28 +15,31 @@ def _create_openmm_virtual_site( interchange: Interchange, - virtual_site: "_VirtualSite", - openff_openmm_particle_map: dict[int | VirtualSiteKey, int], + virtual_site: _VirtualSite | _OpenMMVirtualSite, + openff_openmm_particle_map: dict[int | BaseVirtualSiteKey, int], ) -> openmm.VirtualSite: - # virtual_site.orientations is a list of the _openff_ indices, which is more or less - # the topology index in a topology containing only atoms (no virtual site). This dict, - # _if only looking up atoms_, can be used to map between openff "indices" and - # openmm "indices", where the openff "index" is the atom's index in the (openff) topology - # and the openmm "index" is the atom's index, as a particle, in the openmm system. This - # mapping has a different meaning if looking up a virtual site, but that should not happen here - # as a virtual site's orientation atom should never be a virtual site - openmm_indices: list[int] = [ - openff_openmm_particle_map[openff_index] for openff_index in virtual_site.orientations - ] + if isinstance(virtual_site, _VirtualSite): + # virtual_site.orientations is a list of the _openff_ indices, which is more or less + # the topology index in a topology containing only atoms (no virtual site). This dict, + # _if only looking up atoms_, can be used to map between openff "indices" and + # openmm "indices", where the openff "index" is the atom's index in the (openff) topology + # and the openmm "index" is the atom's index, as a particle, in the openmm system. This + # mapping has a different meaning if looking up a virtual site, but that should not happen here + # as a virtual site's orientation atom should never be a virtual site + openmm_indices: list[int] = [ + openff_openmm_particle_map[openff_index] for openff_index in virtual_site.orientations + ] - # It is assumed that the first "orientation" atom is the "parent" atom. - originwt, xdir, ydir = virtual_site.local_frame_weights - pos = virtual_site.local_frame_positions + # It is assumed that the first "orientation" atom is the "parent" atom. + originwt, xdir, ydir = virtual_site.local_frame_weights + pos = virtual_site.local_frame_positions - return openmm.LocalCoordinatesSite( - openmm_indices, - originwt, - xdir, - ydir, - pos.to_openmm(), - ) + return openmm.LocalCoordinatesSite( + openmm_indices, + originwt, + xdir, + ydir, + pos.to_openmm(), + ) + elif isinstance(virtual_site, _OpenMMVirtualSite): + return virtual_site.to_openmm() diff --git a/openff/interchange/models.py b/openff/interchange/models.py index 4ca9421a0..7b84c062b 100644 --- a/openff/interchange/models.py +++ b/openff/interchange/models.py @@ -332,11 +332,7 @@ def __hash__(self) -> int: return hash((self.this_atom_index, self.other_atom_indices)) -class VirtualSiteKey(TopologyKey): - """ - A unique identifier of a virtual site in the scope of a chemical topology. - """ - +class BaseVirtualSiteKey(TopologyKey): # TODO: Overriding the attribute of a parent class is clumsy, but less grief than # having this not inherit from `TopologyKey`. It might be useful to just have # orientation_atom_indices point to the same thing. @@ -349,19 +345,53 @@ class VirtualSiteKey(TopologyKey): ) type: str = Field(description="The type of this virtual site parameter.") name: str = Field(description="The name of this virtual site parameter.") + + def __hash__(self) -> int: + return hash( + ( + self.orientation_atom_indices, + self.name, + self.type, + ), + ) + + def __repr__(self) -> str: + return f"{self.__class__.__name__} with orientation atom indices {self.orientation_atom_indices}" + + +class ImportedVirtualSiteKey(BaseVirtualSiteKey): + """ + A unique identifier of a virtual site in the scope of a chemical topology. + + Meant to be used with data imported from OpenMM or other engines. + + Use the engine-specific identifier, like `openmm.ThreeParticleAverageSite`, in the "type" field. + """ + + pass + + +class SMIRNOFFVirtualSiteKey(BaseVirtualSiteKey): + """A unique identifier of a SMIRNOFF virtual site in the scope of a chemical topology.""" + match: Literal["once", "all_permutations"] = Field( description="The `match` attribute of the associated virtual site type", ) - def _tuple(self) -> tuple[tuple[int, ...], str, str, str]: - return ( - self.orientation_atom_indices, - self.name, - self.type, - self.match, + def __hash__(self) -> int: + return hash( + ( + self.orientation_atom_indices, + self.name, + self.type, + self.match, + ), ) +VirtualSiteKey = SMIRNOFFVirtualSiteKey + + class PotentialKey(_BaseModel): """ A unique identifier of an instance of physical parameters as applied to a segment of a chemical topology. diff --git a/openff/interchange/smirnoff/_gromacs.py b/openff/interchange/smirnoff/_gromacs.py index 27e876fc7..a09ba0aee 100644 --- a/openff/interchange/smirnoff/_gromacs.py +++ b/openff/interchange/smirnoff/_gromacs.py @@ -36,10 +36,10 @@ RyckaertBellemansDihedral, ) from openff.interchange.models import ( + BaseVirtualSiteKey, BondKey, LibraryChargeTopologyKey, TopologyKey, - VirtualSiteKey, ) MoleculeLike: TypeAlias = Molecule | _SimpleMolecule @@ -89,11 +89,11 @@ def _convert( ] = interchange.topology.identical_molecule_groups virtual_site_molecule_map: dict[ - VirtualSiteKey, + BaseVirtualSiteKey, int, ] = _virtual_site_parent_molecule_mapping(interchange) - molecule_virtual_site_map: dict[int, list[VirtualSiteKey]] = defaultdict(list) + molecule_virtual_site_map: dict[int, list[BaseVirtualSiteKey]] = defaultdict(list) for virtual_site_key, molecule_index in virtual_site_molecule_map.items(): molecule_virtual_site_map[molecule_index].append(virtual_site_key) @@ -140,7 +140,7 @@ def _convert( # when looking up parameters, use the topology index, not the particle index ... # ... or so I think is the expectation of the `TopologyKey`s in the vdW collection topology_index = interchange.topology.atom_index(atom) - key: TopologyKey | VirtualSiteKey | LibraryChargeTopologyKey = TopologyKey( + key: TopologyKey | BaseVirtualSiteKey | LibraryChargeTopologyKey = TopologyKey( atom_indices=(topology_index,), ) @@ -182,13 +182,13 @@ def _convert( epsilon=vdw_parameters["epsilon"].to(unit.kilojoule_per_mole), ) - _partial_charges: dict[int | VirtualSiteKey, float] = dict() + _partial_charges: dict[int | BaseVirtualSiteKey, float] = dict() # Indexed by particle (atom or virtual site) indices for key, charge in interchange["Electrostatics"]._get_charges().items(): if type(key) is TopologyKey: _partial_charges[key.atom_indices[0]] = charge - elif type(key) is VirtualSiteKey: + elif isinstance(key, BaseVirtualSiteKey): _partial_charges[key] = charge else: raise RuntimeError() diff --git a/openff/interchange/smirnoff/_virtual_sites.py b/openff/interchange/smirnoff/_virtual_sites.py index 748bb42ec..7f1955a98 100644 --- a/openff/interchange/smirnoff/_virtual_sites.py +++ b/openff/interchange/smirnoff/_virtual_sites.py @@ -17,7 +17,8 @@ _lookup_virtual_site_parameter, _validated_list_to_array, ) -from openff.interchange.models import PotentialKey, VirtualSiteKey +from openff.interchange.interop._virtual_sites import _ThreeParticleAverageSite +from openff.interchange.models import BaseVirtualSiteKey, ImportedVirtualSiteKey, PotentialKey, SMIRNOFFVirtualSiteKey from openff.interchange.smirnoff._base import SMIRNOFFCollection from openff.interchange.smirnoff._nonbonded import ( SMIRNOFFElectrostaticsCollection, @@ -38,14 +39,14 @@ class SMIRNOFFVirtualSiteCollection(SMIRNOFFCollection): A handler which stores the information necessary to construct virtual sites (virtual particles). """ - key_map: dict[VirtualSiteKey, PotentialKey] = Field( + key_map: dict[BaseVirtualSiteKey, PotentialKey] = Field( dict(), description="A mapping between VirtualSiteKey objects and PotentialKey objects.", ) # type: ignore[assignment] type: Literal["VirtualSites"] = "VirtualSites" expression: Literal[""] = "" - virtual_site_key_topology_index_map: dict[VirtualSiteKey, int] = Field( + virtual_site_key_topology_index_map: dict[BaseVirtualSiteKey, int] = Field( dict(), description="A mapping between VirtualSiteKey objects (stored analogously to TopologyKey objects" "in other handlers) and topology indices describing the associated virtual site", @@ -107,7 +108,7 @@ def store_matches( for orientation in orientations: orientation_indices = orientation.topology_atom_indices - virtual_site_key = VirtualSiteKey( + virtual_site_key = SMIRNOFFVirtualSiteKey( parent_atom_index=parent_index, orientation_atom_indices=orientation_indices, type=parameter.type, @@ -338,40 +339,57 @@ def local_frame_coordinates(self) -> Quantity: def _create_virtual_site_object( - virtual_site_key: VirtualSiteKey, + virtual_site_key: BaseVirtualSiteKey, virtual_site_potential, # interchange: "Interchange", # non_bonded_force: openmm.NonbondedForce, ) -> _VirtualSite: orientations = virtual_site_key.orientation_atom_indices - if virtual_site_key.type == "BondCharge": - return _BondChargeVirtualSite( - type="BondCharge", - distance=virtual_site_potential.parameters["distance"], - orientations=orientations, - ) - elif virtual_site_key.type == "MonovalentLonePair": - return _MonovalentLonePairVirtualSite( - type="MonovalentLonePair", - distance=virtual_site_potential.parameters["distance"], - out_of_plane_angle=virtual_site_potential.parameters["outOfPlaneAngle"], - in_plane_angle=virtual_site_potential.parameters["inPlaneAngle"], - orientations=orientations, - ) - elif virtual_site_key.type == "DivalentLonePair": - return _DivalentLonePairVirtualSite( - type="DivalentLonePair", - distance=virtual_site_potential.parameters["distance"], - out_of_plane_angle=virtual_site_potential.parameters["outOfPlaneAngle"], - orientations=orientations, - ) - elif virtual_site_key.type == "TrivalentLonePair": - return _TrivalentLonePairVirtualSite( - type="TrivalentLonePair", - distance=virtual_site_potential.parameters["distance"], - orientations=orientations, - ) + # this check is basically "is this a SMIRNOFF virtual site?" + # see commment at its class definition + if isinstance(virtual_site_key, SMIRNOFFVirtualSiteKey): + if virtual_site_key.type == "BondCharge": + return _BondChargeVirtualSite( + type="BondCharge", + distance=virtual_site_potential.parameters["distance"], + orientations=orientations, + ) + elif virtual_site_key.type == "MonovalentLonePair": + return _MonovalentLonePairVirtualSite( + type="MonovalentLonePair", + distance=virtual_site_potential.parameters["distance"], + out_of_plane_angle=virtual_site_potential.parameters["outOfPlaneAngle"], + in_plane_angle=virtual_site_potential.parameters["inPlaneAngle"], + orientations=orientations, + ) + elif virtual_site_key.type == "DivalentLonePair": + return _DivalentLonePairVirtualSite( + type="DivalentLonePair", + distance=virtual_site_potential.parameters["distance"], + out_of_plane_angle=virtual_site_potential.parameters["outOfPlaneAngle"], + orientations=orientations, + ) + elif virtual_site_key.type == "TrivalentLonePair": + return _TrivalentLonePairVirtualSite( + type="TrivalentLonePair", + distance=virtual_site_potential.parameters["distance"], + orientations=orientations, + ) + + else: + raise NotImplementedError(virtual_site_key.type) + + # TODO: This case shouldn't be in openff/interchange/smirnoff! + elif isinstance(virtual_site_key, ImportedVirtualSiteKey): + if virtual_site_key.type == "ThreeParticleAverageSite": + return _ThreeParticleAverageSite( # type: ignore[return-value] + particles=virtual_site_key.orientation_atom_indices, + weights=virtual_site_potential.parameters["weights"], + ) + + else: + raise NotImplementedError(virtual_site_key.type) else: raise NotImplementedError(virtual_site_key.type) @@ -464,28 +482,44 @@ def _generate_positions( conformer: _Quantity | None = None, ) -> Quantity: # TODO: Capture these objects instead of generating them on-the-fly so many times + try: + local_frame_coordinates = numpy.vstack( + [ + _create_virtual_site_object( + virtual_site_key, + virtual_site_collection.potentials[potential_key], + ).local_frame_coordinates + for virtual_site_key, potential_key in virtual_site_collection.key_map.items() + ], + ) - local_frame_coordinates = numpy.vstack( - [ - _create_virtual_site_object( - virtual_site_key, - virtual_site_collection.potentials[potential_key], - ).local_frame_coordinates - for virtual_site_key, potential_key in virtual_site_collection.key_map.items() - ], - ) + local_coordinate_frames = _build_local_coordinate_frames( + interchange, + virtual_site_collection, + ) - local_coordinate_frames = _build_local_coordinate_frames( - interchange, - virtual_site_collection, - ) + virtual_site_positions = _convert_local_coordinates( + local_frame_coordinates, + local_coordinate_frames, + ) - virtual_site_positions = _convert_local_coordinates( - local_frame_coordinates, - local_coordinate_frames, - ) + return Quantity( + virtual_site_positions, + unit.nanometer, + ) + except AttributeError: + # Handle case of virtual sites imported from OpenMM + # TODO: Handle case of mixed SMIRNOFF and OpenMM virtual sites + _virtual_site_positions: list[Quantity] = list() - return Quantity( - virtual_site_positions, - unit.nanometer, - ) + for key, val in virtual_site_collection.key_map.items(): + atom_indices = key.orientation_atom_indices + weights = virtual_site_collection.potentials[val].parameters["weights"].m + + assert len(atom_indices) == len(weights) + + _virtual_site_positions.append( + sum(interchange.positions[atom_indices[i]] * weights[i] for i in range(len(weights))), + ) + + return numpy.vstack(_virtual_site_positions)