Add support for None bond anchors

Add test for single topology end-states

WIP

Add a test for multiple disconnected components.

Fix typo in kwarg.

Author: proteneer

tests/test_dummy.py

@@ -1,7 +1,11 @@
 import pytest
 from rdkit import Chem
 
-from timemachine.fe.dummy import generate_dummy_group_assignments
+from timemachine.fe.dummy import (
+    ZeroBondAnchorWarning,
+    generate_anchored_dummy_group_assignments,
+    generate_dummy_group_assignments,
+)
 from timemachine.graph_utils import convert_to_nx
 
 # These tests check the various utilities used to turn off interactions
@@ -92,4 +96,86 @@ def test_generate_dummy_group_assignments():
 def test_generate_dummy_group_assignments_empty_core():
     g = convert_to_nx(Chem.MolFromSmiles("OC1COO1"))
     core = []
-    assert list(generate_dummy_group_assignments(g, core)) == []
+
+    with pytest.warns(ZeroBondAnchorWarning):
+        dgas = list(generate_dummy_group_assignments(g, core))
+
+    assert equivalent_assignment(dgas, [{None: {0, 1, 2, 3, 4}}])
+
+
+def test_generate_dummy_group_assignments_full_core():
+    g = convert_to_nx(Chem.MolFromSmiles("OC1COO1"))
+    core = [0, 1, 2, 3, 4]
+    dgas = list(generate_dummy_group_assignments(g, core))
+    assert equivalent_assignment(dgas, [{}])
+
+
+def test_generate_angle_anchor_dummy_group_assignments():
+    # Test that if we break a core-core bond, we only have one valid
+    # choice of the angle anchor
+    #
+    #      O0          O0
+    #      |           |
+    #      C1          C1
+    #     / \           \
+    #    O4  C2  ->  O4  C2
+    #     \ /         \ /
+    #      O3          O3
+    g_a = convert_to_nx(Chem.MolFromSmiles("OC1COO1"))
+    g_b = convert_to_nx(Chem.MolFromSmiles("OCCOO"))
+    core_a = [1, 2, 3, 4]
+    core_b = [1, 2, 3, 4]
+
+    dgas = list(generate_dummy_group_assignments(g_a, core_a))
+    expected_dga = {1: {0}}
+
+    assert equivalent_assignment(dgas, [expected_dga])
+
+    # forward direction
+    anchored_dummy_group_assignments = generate_anchored_dummy_group_assignments(expected_dga, g_a, g_b, core_a, core_b)
+
+    anchored_dummy_group_assignments = list(anchored_dummy_group_assignments)
+    assert len(anchored_dummy_group_assignments) == 1
+    assert anchored_dummy_group_assignments[0] == {1: (2, frozenset({0}))}
+
+    # reverse direction
+    anchored_dummy_group_assignments = generate_anchored_dummy_group_assignments(expected_dga, g_b, g_a, core_b, core_a)
+
+    anchored_dummy_group_assignments = list(anchored_dummy_group_assignments)
+    assert len(anchored_dummy_group_assignments) == 1
+    assert anchored_dummy_group_assignments[0] == {1: (2, frozenset({0}))}
+
+    # Test that providing an empty core and a None bond anchor results in a None angle anchor
+    anchored_dummy_group_assignments = generate_anchored_dummy_group_assignments(
+        {None: {0, 1, 2, 3, 4}}, g_a, g_b, core_atoms_a=[], core_atoms_b=[]
+    )
+
+    anchored_dummy_group_assignments = list(anchored_dummy_group_assignments)
+
+    assert anchored_dummy_group_assignments == [{None: (None, frozenset({0, 1, 2, 3, 4}))}]
+
+
+def test_multiple_none_dummy_groups():
+    # Test that if we break a core-core bond, we only have one valid
+    # choice of the angle anchor
+    #
+    # mol_a: H-O-H.H-O-H
+    # mol_b: H-O-H.H-O-H
+    g_a = convert_to_nx(Chem.AddHs(Chem.MolFromSmiles("O.O")))
+    g_b = convert_to_nx(Chem.AddHs(Chem.MolFromSmiles("O.O")))
+    core_a = []
+    core_b = []
+
+    dgas = list(generate_dummy_group_assignments(g_a, core_a, assert_single_connected_component=False))
+    expected_dga = {None: {0, 1, 2, 3, 4, 5}}
+
+    assert equivalent_assignment(dgas, [expected_dga])
+
+    # Test that providing an empty core and a None bond anchor results in a None angle anchor
+    anchored_dummy_group_assignments = generate_anchored_dummy_group_assignments(
+        {None: {0, 1, 2, 3, 4, 5}}, g_a, g_b, core_a, core_b
+    )
+
+    anchored_dummy_group_assignments = list(anchored_dummy_group_assignments)
+
+    assert anchored_dummy_group_assignments == [{None: (None, frozenset({0, 1, 2, 3, 4, 5}))}]

tests/test_single_topology.py

@@ -20,7 +20,7 @@
     NBParamIdx,
 )
 from timemachine.fe import atom_mapping, single_topology
-from timemachine.fe.dummy import MultipleAnchorWarning, canonicalize_bond
+from timemachine.fe.dummy import MultipleBondAnchorWarning, canonicalize_bond
 from timemachine.fe.interpolate import align_nonbonded_idxs_and_params, linear_interpolation
 from timemachine.fe.single_topology import (
     AtomMapMixin,
@@ -36,6 +36,7 @@
     setup_dummy_interactions_from_ff,
 )
 from timemachine.fe.system import minimize_scipy, simulate_system
+from timemachine.fe.topology import BaseTopology
 from timemachine.fe.utils import get_mol_name, get_romol_conf, read_sdf, read_sdf_mols_by_name, set_mol_name
 from timemachine.ff import Forcefield
 from timemachine.md import minimizer
@@ -330,7 +331,7 @@ def test_find_dummy_groups_and_multiple_anchors():
 
     core_pairs = np.array([[1, 1], [2, 2]])
 
-    with pytest.warns(MultipleAnchorWarning):
+    with pytest.warns(MultipleBondAnchorWarning):
         dgs = single_topology.find_dummy_groups_and_anchors(mol_a, mol_b, core_pairs[:, 0], core_pairs[:, 1])
         assert dgs == {1: (2, {0})} or dgs == {2: (1, {0})}
 
@@ -353,7 +354,7 @@ def test_find_dummy_groups_and_multiple_anchors():
     core_a = [0, 1, 2, 3]
     core_b = [2, 1, 4, 3]
 
-    with pytest.warns(MultipleAnchorWarning):
+    with pytest.warns(MultipleBondAnchorWarning):
         dgs = single_topology.find_dummy_groups_and_anchors(mol_a, mol_b, core_a, core_b)
         assert dgs == {1: (2, {0})}
 
@@ -1743,3 +1744,48 @@ def test_hif2a_end_state_symmetry_nightly_test():
         print("testing", mol_a.GetProp("_Name"), "->", mol_b.GetProp("_Name"))
         core = atom_mapping.get_cores(mol_a, mol_b, **DEFAULT_ATOM_MAPPING_KWARGS)[0]
         assert_symmetric_interpolation(mol_a, mol_b, core)
+
+
+def test_empty_core():
+    # test that an empty core results in an end-state with identical energies and forces to independent molecules
+    with path_to_internal_file("timemachine.testsystems.data", "ligands_40.sdf") as path_to_ligand:
+        mols = read_sdf(path_to_ligand)
+
+    mol_a = mols[0]
+    mol_b = mols[1]
+
+    x_a = get_romol_conf(mol_a)
+    x_b = get_romol_conf(mol_b)
+    core = np.zeros((0, 2))
+
+    ff = Forcefield.load_default()
+    st = SingleTopology(mol_a, mol_b, core, ff)
+    lhs = st.setup_intermediate_state(0.0)
+    rhs = st.setup_intermediate_state(1.0)
+
+    x_0 = st.combine_confs(x_a, x_b, 0.0)
+    x_1 = st.combine_confs(x_a, x_b, 1.0)
+
+    np.testing.assert_array_equal(x_0, x_1)
+
+    # lhs and rhs do not have identical total energies since dummy molecule is in a softened state
+    lhs_U_fn = lhs.get_U_fn()
+    rhs_U_fn = rhs.get_U_fn()
+
+    assert lhs_U_fn(x_0) != rhs_U_fn(x_0)
+
+    # however, the bond, angle, improper energies should be bitwise identical
+    assert lhs.bond(x_0, None) == rhs.bond(x_0, None)
+    assert lhs.angle(x_0, None) == rhs.angle(x_0, None)
+    assert lhs.improper(x_0, None) == rhs.improper(x_0, None)
+
+    # we should also be consistent with vanilla base topologies
+    bt_a = BaseTopology(mol_a, ff)
+    bt_b = BaseTopology(mol_b, ff)
+
+    ref_a = bt_a.setup_end_state()
+    ref_b = bt_b.setup_end_state()
+
+    np.testing.assert_almost_equal(lhs.bond(x_0, None), ref_a.bond(x_a, None) + ref_b.bond(x_b, None))
+    np.testing.assert_almost_equal(lhs.angle(x_0, None), ref_a.angle(x_a, None) + ref_b.angle(x_b, None))
+    np.testing.assert_almost_equal(lhs.improper(x_0, None), ref_a.improper(x_a, None) + ref_b.improper(x_b, None))

timemachine/fe/dummy.py

@@ -7,14 +7,18 @@
 import networkx as nx
 
 
-class MultipleAnchorWarning(UserWarning):
+class MultipleBondAnchorWarning(UserWarning):
+    pass
+
+
+class ZeroBondAnchorWarning(UserWarning):
     pass
 
 
 def generate_dummy_group_assignments(
-    bond_graph: nx.Graph, core_atoms: Collection[int]
+    bond_graph: nx.Graph, core_atoms: Collection[int], assert_single_connected_component: bool = True
 ) -> Iterator[dict[int, frozenset[int]]]:
-    """Returns an iterator over dummy group assignments (i.e., candidate partitionings of dummy atoms with each
+    """Returns an iterator over all possible dummy group assignments (i.e., candidate partitionings of dummy atoms with each
     partition assigned a bond anchor atom) for a given molecule (represented as a bond graph) and set of core atoms.
 
     A dummy group is a set of dummy atoms that are inserted or deleted in alchemical free energy calculations. The
@@ -64,18 +68,26 @@ def generate_dummy_group_assignments(
        latter might be desirable for efficiency but is more complicated to implement).
     """
     assert len(set(core_atoms)) == len(core_atoms)
-    assert len(list(nx.connected_components(bond_graph))) == 1
+
+    if assert_single_connected_component:
+        assert len(list(nx.connected_components(bond_graph))) == 1
 
     core_atoms_ = frozenset(core_atoms)
     dummy_atoms = frozenset(bond_graph.nodes()) - core_atoms_
     induced_g = nx.subgraph(bond_graph, dummy_atoms)
 
     def get_bond_anchors(dummy_group):
         bond_anchors = {n for dummy_atom in dummy_group for n in bond_graph.neighbors(dummy_atom) if n in core_atoms}
-        if len(bond_anchors) > 1:
+        if len(bond_anchors) == 0:
+            warnings.warn(
+                f"No bond anchors found for dummy group: {dummy_group}",
+                ZeroBondAnchorWarning,
+            )
+            bond_anchors = set([None])
+        elif len(bond_anchors) > 1:
             warnings.warn(
                 f"Multiple bond anchors {bond_anchors} found for dummy group: {dummy_group}",
-                MultipleAnchorWarning,
+                MultipleBondAnchorWarning,
             )
         return bond_anchors
 
@@ -138,11 +150,16 @@ def generate_anchored_dummy_group_assignments(
         each element is a mapping from bond anchor atom to the pair (angle anchor atom, dummy group)
     """
 
+    assert len(core_atoms_a) == len(core_atoms_b)
+
     core_bonds_c = get_core_bonds(bond_graph_a.edges(), bond_graph_b.edges(), core_atoms_a, core_atoms_b)
     c_to_b = {c: b for c, b in enumerate(core_atoms_b)}
     core_bonds_b = frozenset(translate_bonds(core_bonds_c, c_to_b))
 
     def get_angle_anchors(bond_anchor):
+        if bond_anchor is None:
+            return [None]
+
         valid_angle_anchors = [
             angle_anchor
             for angle_anchor in [n for n in bond_graph_b.neighbors(bond_anchor) if n in core_atoms_b]

timemachine/fe/single_topology.py

@@ -144,7 +144,7 @@ def setup_dummy_bond_and_chiral_interactions(
     root_anchor_atom: int,
     core_atoms: NDArray,
 ):
-    assert root_anchor_atom in core_atoms
+    assert root_anchor_atom is None or root_anchor_atom in core_atoms
 
     dummy_group_arr = np.array(list(dummy_group))
 
@@ -286,7 +286,7 @@ def setup_dummy_interactions(
     (bonded_idxs, bonded_params)
         Returns bonds, angles, and improper idxs and parameters.
     """
-    assert root_anchor_atom in core_atoms
+    assert root_anchor_atom is None or root_anchor_atom in core_atoms
 
     dummy_angle_idxs = []
     dummy_angle_params = []
@@ -663,9 +663,8 @@ def concatenate(arrays, empty_shape, empty_dtype):
     )
 
     num_atoms = mol_a.GetNumAtoms() + mol_b.GetNumAtoms() - len(core)
-    assert get_num_connected_components(num_atoms, bond_potential.potential.idxs) == 1, (
-        "hybrid molecule has multiple connected components"
-    )
+    if get_num_connected_components(num_atoms, bond_potential.potential.idxs) > 1:
+        warnings.warn("Hybrid molecule has multiple connected components")
 
     return GuestSystem(
         bond=bond_potential,