rotational bond keys update test fix

avcopan · Jun 25, 2024 · c125651 · c125651
1 parent e398872
commit c125651
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 5 deletions.
diff --git a/automol/graph/_0embed.py b/automol/graph/_0embed.py
@@ -330,7 +330,7 @@ def distance_bounds_matrices(
     # 3. reopen bounds on the torsions from the reactant
     # (also opens >5-atom distances)
     if relax_torsions:
-        rot_bkeys = rotational_bond_keys(gra)
+        rot_bkeys = rotational_bond_keys(gra, extend_lin_seg=False)
 
         # open up torsions at rotational bonds
         tors_ijs = [

diff --git a/automol/graph/base/_06heur.py b/automol/graph/base/_06heur.py
@@ -280,6 +280,7 @@ def rotational_coordinates(
     lin_keys: Optional[List[int]] = None,
     with_h_rotors: bool = True,
     with_ch_rotors: bool = True,
+    extend_lin_seg: bool = False
 ):
     """Get torsion coordinates for rotational segments
 
@@ -308,6 +309,7 @@ def rotational_coordinates(
         lin_keys=lin_keys,
         with_h_rotors=with_h_rotors,
         with_ch_rotors=with_ch_rotors,
+        extend_lin_seg=extend_lin_seg
     )
 
     coo_keys = []

diff --git a/automol/tests/test_graph.py b/automol/tests/test_graph.py
@@ -1566,7 +1566,7 @@ def test__rotational_bond_keys():
 
     # Check that we don't misidentify rotational bond keys
     cgr = automol.smiles.graph("C#CC=C")
-    assert graph.rotational_bond_keys(cgr) == frozenset()
+    assert graph.rotational_bond_keys(cgr, extend_lin_seg) == frozenset()
 
 
 def test__rotational_segment_keys():
@@ -1615,7 +1615,7 @@ def test__rotational_coordinates():
     tors_dct = automol.zmat.torsion_coordinates(zma, zgra)
     coo_key_lst0 = set(map(tuple, map(reversed, tors_dct.values())))
 
-    coo_key_lst = graph.rotational_coordinates(zgra, segment=False)
+    coo_key_lst = graph.rotational_coordinates(zgra, extend_lin_seg=True, segment=False)
     assert coo_key_lst == coo_key_lst0, f"{coo_key_lst} != {coo_key_lst0}"
 
     # Make sure the segment version runs

diff --git a/automol/tests/test_graph_ts.py b/automol/tests/test_graph_ts.py
@@ -1157,8 +1157,8 @@ def test__rotational_bond_keys():
         },
     )
     rtsg = graph.ts.reverse(tsg)
-    assert graph.rotational_bond_keys(tsg) == frozenset()
-    assert graph.rotational_bond_keys(rtsg) == frozenset()
+    assert graph.rotational_bond_keys(tsg, extend_lin_seg=False) == frozenset()
+    assert graph.rotational_bond_keys(rtsg, extend_lin_seg=False) == frozenset()
 
 
 def test__radical_atom_keys():