Don't include constrained bonds / angles from OpenFF (#116)

smee/converters/openff/_openff.py

@@ -269,6 +269,7 @@ def convert_handlers(
     potentials: (
         list[tuple[smee.TensorPotential, list[smee.ParameterMap]]] | None
     ) = None,
+    constraints: list[smee.TensorConstraints | None] | None = None,
 ) -> list[tuple[smee.TensorPotential, list[smee.ParameterMap]]]:
     """Convert a set of SMIRNOFF parameter handlers into a set of tensor potentials.
 
@@ -277,6 +278,7 @@ def convert_handlers(
             objects to convert.
         topologies: The topologies associated with each interchange object.
         v_site_maps: The v-site maps associated with each interchange object.
+        constraints: Any distance constraints between atoms.
         potentials: Already converted parameter handlers that may be required as
             dependencies.
 
@@ -315,6 +317,8 @@ def convert_handlers(
     if handler_type not in _CONVERTERS:
         raise NotImplementedError(f"{handler_type} handlers is not yet supported.")
 
+    constraints = [None] * len(topologies) if constraints is None else constraints
+
     converter = _CONVERTERS[handler_type]
     converter_spec = inspect.signature(converter.fn)
     converter_kwargs = {}
@@ -324,6 +328,11 @@ def convert_handlers(
     if "v_site_maps" in converter_spec.parameters:
         assert v_site_maps is not None, "v-site maps must be provided"
         converter_kwargs["v_site_maps"] = v_site_maps
+    if "constraints" in converter_spec.parameters:
+        constraint_idxs = [[] if v is None else v.idxs.tolist() for v in constraints]
+        unique_idxs = [{tuple(sorted(idxs)) for idxs in v} for v in constraint_idxs]
+
+        converter_kwargs["constraints"] = unique_idxs
 
     potentials_by_type = (
         {}
@@ -508,7 +517,9 @@ def convert_interchange(
         ):
             continue
 
-        converted = convert_handlers(handlers, topologies, v_site_maps, converted)
+        converted = convert_handlers(
+            handlers, topologies, v_site_maps, converted, constraints
+        )
 
     # handlers may either return multiple potentials, or condense multiple already
     # converted potentials into a single one (e.g. electrostatics into some polarizable

smee/converters/openff/valence.py

@@ -12,6 +12,67 @@
 _KCAL_PER_MOL = openff.units.unit.kilocalories / openff.units.unit.mole
 
 
+def strip_constrained_bonds(
+    parameter_maps: list[smee.ValenceParameterMap],
+    constraints: list[set[tuple[int, int]]],
+):
+    """Remove bonded interactions between distance-constrained atoms.
+
+    Args:
+        parameter_maps: The parameter maps to strip.
+        constraints: The distanced constrained bonds to exclude for each parameter map.
+    """
+
+    for parameter_map, bonds_to_exclude in zip(
+        parameter_maps, constraints, strict=True
+    ):
+        bonds_to_exclude = {tuple(sorted(idxs)) for idxs in bonds_to_exclude}
+
+        bond_idxs = [
+            tuple(sorted(idxs)) for idxs in parameter_map.particle_idxs.tolist()
+        ]
+        include = [idxs not in bonds_to_exclude for idxs in bond_idxs]
+
+        parameter_map.particle_idxs = parameter_map.particle_idxs[include]
+        parameter_map.assignment_matrix = parameter_map.assignment_matrix.to_dense()[
+            include, :
+        ].to_sparse()
+
+
+def strip_constrained_angles(
+    parameter_maps: list[smee.ValenceParameterMap],
+    constraints: list[set[tuple[int, int]]],
+):
+    """Remove angle interactions between angles where all three atoms are constrained
+    with distance constraints.
+
+    Args:
+        parameter_maps: The parameter maps to strip.
+        constraints: The distanced constrained bonds to exclude for each parameter map.
+    """
+
+    def is_constrained(idxs_, excluded):
+        bonds = {
+            tuple(sorted([idxs_[0], idxs_[1]])),
+            tuple(sorted([idxs_[0], idxs_[2]])),
+            tuple(sorted([idxs_[1], idxs_[2]])),
+        }
+        return len(bonds & excluded) == 3
+
+    for parameter_map, bonds_to_exclude in zip(
+        parameter_maps, constraints, strict=True
+    ):
+        bonds_to_exclude = {tuple(sorted(idxs)) for idxs in bonds_to_exclude}
+
+        angle_idxs = parameter_map.particle_idxs.tolist()
+        include = [not is_constrained(idxs, bonds_to_exclude) for idxs in angle_idxs]
+
+        parameter_map.particle_idxs = parameter_map.particle_idxs[include]
+        parameter_map.assignment_matrix = parameter_map.assignment_matrix.to_dense()[
+            include, :
+        ].to_sparse()
+
+
 def convert_valence_handlers(
     handlers: list[openff.interchange.smirnoff.SMIRNOFFCollection],
     handler_type: str,
@@ -64,17 +125,29 @@ def convert_valence_handlers(
 )
 def convert_bonds(
     handlers: list[openff.interchange.smirnoff.SMIRNOFFBondCollection],
+    constraints: list[set[tuple[int, int]]],
 ) -> tuple[smee.TensorPotential, list[smee.ValenceParameterMap]]:
-    return convert_valence_handlers(handlers, "Bonds", ("k", "length"))
+    potential, parameter_maps = convert_valence_handlers(
+        handlers, "Bonds", ("k", "length")
+    )
+    strip_constrained_bonds(parameter_maps, constraints)
+
+    return potential, parameter_maps
 
 
 @smee.converters.smirnoff_parameter_converter(
     "Angles", {"k": _KCAL_PER_MOL / _RADIANS**2, "angle": _RADIANS}
 )
 def convert_angles(
     handlers: list[openff.interchange.smirnoff.SMIRNOFFAngleCollection],
+    constraints: list[set[tuple[int, int]]],
 ) -> tuple[smee.TensorPotential, list[smee.ValenceParameterMap]]:
-    return convert_valence_handlers(handlers, "Angles", ("k", "angle"))
+    potential, parameter_maps = convert_valence_handlers(
+        handlers, "Angles", ("k", "angle")
+    )
+    strip_constrained_angles(parameter_maps, constraints)
+
+    return potential, parameter_maps
 
 
 @smee.converters.smirnoff_parameter_converter(

smee/tests/convertors/openff/test_valence.py

@@ -1,3 +1,7 @@
+import openff.interchange
+import openff.toolkit
+import pytest
+
 from smee.converters.openff.valence import (
     convert_angles,
     convert_bonds,
@@ -9,7 +13,7 @@
 def test_convert_bonds(ethanol, ethanol_interchange):
     bond_collection = ethanol_interchange.collections["Bonds"]
 
-    potential, parameter_maps = convert_bonds([bond_collection])
+    potential, parameter_maps = convert_bonds([bond_collection], [set()])
 
     assert potential.type == "Bonds"
     assert potential.fn == "k/2*(r-length)**2"
@@ -56,14 +60,125 @@ def test_convert_bonds(ethanol, ethanol_interchange):
     assert actual_parameters == expected_parameters
 
 
-def test_convert_angles(ethanol, ethanol_interchange):
+def test_convert_bonds_with_constraints(ethanol):
+    interchange = openff.interchange.Interchange.from_smirnoff(
+        openff.toolkit.ForceField("openff-1.3.0.offxml"), ethanol.to_topology()
+    )
+
+    bond_collection = interchange.collections["Bonds"]
+
+    constraints = {
+        (bond.atom1_index, bond.atom2_index)
+        for bond in ethanol.bonds
+        if bond.atom1.atomic_number == 1 or bond.atom2.atomic_number == 1
+    }
+
+    potential, [parameter_map] = convert_bonds([bond_collection], [constraints])
+    parameter_keys = [key.id for key in potential.parameter_keys]
+
+    assert len(parameter_map.assignment_matrix) == len(parameter_map.particle_idxs)
+    assignment_matrix = parameter_map.assignment_matrix.to_dense()
+
+    actual_parameters = {
+        tuple(particle_idxs.tolist()): parameter_keys[parameter_idxs.nonzero()]
+        for parameter_idxs, particle_idxs in zip(
+            assignment_matrix, parameter_map.particle_idxs, strict=True
+        )
+    }
+    expected_parameters = {(0, 2): "[#6:1]-[#8:2]", (1, 2): "[#6X4:1]-[#6X4:2]"}
+
+    assert actual_parameters == expected_parameters
+
+
+@pytest.mark.parametrize("with_constraints", [True, False])
+def test_convert_angles_etoh(ethanol, ethanol_interchange, with_constraints):
     angle_collection = ethanol_interchange.collections["Angles"]
 
-    potential, parameter_maps = convert_angles([angle_collection])
+    h_bond_idxs = {
+        (bond.atom1_index, bond.atom2_index)
+        for bond in ethanol.bonds
+        if bond.atom1.atomic_number == 1 or bond.atom2.atomic_number == 1
+    }
+    constraints = set() if not with_constraints else h_bond_idxs
+
+    potential, parameter_maps = convert_angles([angle_collection], [constraints])
 
     assert potential.type == "Angles"
     assert potential.fn == "k/2*(theta-angle)**2"
 
+    assert potential.attributes is None
+    assert potential.attribute_cols is None
+
+    assert potential.parameter_cols == ("k", "angle")
+
+    parameter_keys = [key.id for key in potential.parameter_keys]
+    expected_parameter_keys = [
+        "[#1:1]-[#6X4:2]-[#1:3]",
+        "[*:1]-[#8:2]-[*:3]",
+        "[*:1]~[#6X4:2]-[*:3]",
+    ]
+    assert sorted(parameter_keys) == sorted(expected_parameter_keys)
+
+    assert potential.parameters.shape == (3, 2)
+
+    assert len(parameter_maps) == 1
+    parameter_map = parameter_maps[0]
+
+    assert len(parameter_map.assignment_matrix) == len(parameter_map.particle_idxs)
+    assignment_matrix = parameter_map.assignment_matrix.to_dense()
+
+    actual_parameters = {
+        tuple(particle_idxs.tolist()): parameter_keys[parameter_idxs.nonzero()]
+        for parameter_idxs, particle_idxs in zip(
+            assignment_matrix, parameter_map.particle_idxs, strict=True
+        )
+    }
+    expected_parameters = {
+        (0, 2, 1): "[*:1]~[#6X4:2]-[*:3]",
+        (0, 2, 7): "[*:1]~[#6X4:2]-[*:3]",
+        (0, 2, 8): "[*:1]~[#6X4:2]-[*:3]",
+        (1, 2, 7): "[*:1]~[#6X4:2]-[*:3]",
+        (1, 2, 8): "[*:1]~[#6X4:2]-[*:3]",
+        (2, 0, 3): "[*:1]-[#8:2]-[*:3]",
+        (2, 1, 4): "[*:1]~[#6X4:2]-[*:3]",
+        (2, 1, 5): "[*:1]~[#6X4:2]-[*:3]",
+        (2, 1, 6): "[*:1]~[#6X4:2]-[*:3]",
+        (4, 1, 5): "[#1:1]-[#6X4:2]-[#1:3]",
+        (4, 1, 6): "[#1:1]-[#6X4:2]-[#1:3]",
+        (5, 1, 6): "[#1:1]-[#6X4:2]-[#1:3]",
+        (7, 2, 8): "[#1:1]-[#6X4:2]-[#1:3]",
+    }
+
+    assert actual_parameters == expected_parameters
+
+
+@pytest.mark.parametrize("with_constraints", [True, False])
+def test_convert_angle_water(with_constraints):
+    interchange = openff.interchange.Interchange.from_smirnoff(
+        openff.toolkit.ForceField("openff-1.3.0.offxml"),
+        openff.toolkit.Molecule.from_mapped_smiles("[O:1]([H:2])[H:3]").to_topology(),
+    )
+
+    angle_collection = interchange.collections["Angles"]
+
+    constraints = set() if not with_constraints else {(0, 1), (0, 2), (1, 2)}
+
+    potential, [parameter_map] = convert_angles([angle_collection], [constraints])
+    parameter_keys = [key.id for key in potential.parameter_keys]
+
+    assert len(parameter_map.assignment_matrix) == len(parameter_map.particle_idxs)
+    assignment_matrix = parameter_map.assignment_matrix.to_dense()
+
+    actual_parameters = {
+        tuple(particle_idxs.tolist()): parameter_keys[parameter_idxs.nonzero()]
+        for parameter_idxs, particle_idxs in zip(
+            assignment_matrix, parameter_map.particle_idxs, strict=True
+        )
+    }
+    expected_parameters = {} if with_constraints else {(1, 0, 2): "[*:1]-[#8:2]-[*:3]"}
+
+    assert actual_parameters == expected_parameters
+
 
 def test_convert_propers(ethanol, ethanol_interchange):
     proper_collection = ethanol_interchange.collections["ProperTorsions"]

smee/tests/convertors/test_openmm.py

@@ -50,7 +50,7 @@ def _compare_smee_and_interchange(
 ):
     system_smee = convert_to_openmm_system(tensor_ff, tensor_system)
     assert isinstance(system_smee, openmm.System)
-    system_interchange = interchange.to_openmm(False, True)
+    system_interchange = interchange.to_openmm(False, False)
 
     coords += (numpy.random.randn(*coords.shape) * 0.1) * openmm.unit.angstrom
 
@@ -145,25 +145,36 @@ def compare_vec3(a: openmm.Vec3, b: openmm.Vec3):
         )
 
 
-def test_convert_to_openmm_system_vacuum():
+@pytest.mark.parametrize("with_constraints", [True, False])
+def test_convert_to_openmm_system_vacuum(with_constraints):
     # carbonic acid has impropers, 1-5 interactions so should test most convertors
     mol = openff.toolkit.Molecule.from_smiles("OC(=O)O")
     mol.generate_conformers(n_conformers=1)
 
     coords = mol.conformers[0].m_as(openff.units.unit.angstrom)
     coords = coords * openmm.unit.angstrom
 
+    force_field = openff.toolkit.ForceField(
+        "openff-2.0.0.offxml"
+        if with_constraints
+        else "openff_unconstrained-2.0.0.offxml"
+    )
     interchange = openff.interchange.Interchange.from_smirnoff(
-        openff.toolkit.ForceField("openff-2.0.0.offxml"), mol.to_topology()
+        force_field, mol.to_topology()
     )
 
     tensor_ff, [tensor_top] = smee.converters.convert_interchange(interchange)
 
     _compare_smee_and_interchange(tensor_ff, tensor_top, interchange, coords, None)
 
 
-def test_convert_to_openmm_system_periodic():
-    ff = openff.toolkit.ForceField("openff-2.0.0.offxml")
+@pytest.mark.parametrize("with_constraints", [True, False])
+def test_convert_to_openmm_system_periodic(with_constraints):
+    ff = openff.toolkit.ForceField(
+        "openff-2.0.0.offxml"
+        if with_constraints
+        else "openff_unconstrained-2.0.0.offxml"
+    )
     top = openff.toolkit.Topology()
 
     interchanges = []