tests

Author: skfegan

CodeEntropy/entropy/nodes/vibrational.py

@@ -109,7 +109,7 @@ def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> Dict[str, Any]
                         residues=rep_mol.residues,
                         force_ua=force_cov["ua"],
                         torque_ua=torque_cov["ua"],
-                        flexible=flexible["ua"],
+                        flexible_ua=flexible["ua"],
                         ua_frame_counts=ua_frame_counts,
                         reporter=reporter,
                         n_frames_default=shared_data.get("n_frames", 0),
@@ -224,7 +224,7 @@ def _compute_united_atom_entropy(
         residues: Any,
         force_ua: Mapping[CovKey, Any],
         torque_ua: Mapping[CovKey, Any],
-        flexible: Any,
+        flexible_ua: Any,
         ua_frame_counts: Mapping[CovKey, int],
         reporter: Optional[Any],
         n_frames_default: int,
@@ -259,7 +259,7 @@ def _compute_united_atom_entropy(
             key = (group_id, res_id)
             fmat = force_ua.get(key)
             tmat = torque_ua.get(key)
-            flexible = flexible[key]
+            flexible = flexible_ua.get(key)
 
             pair = self._compute_force_torque_entropy(
                 ve=ve,

CodeEntropy/entropy/vibrational.py

@@ -123,9 +123,12 @@ def _entropy_components(
             Array of entropy components (J/mol/K) for each valid mode.
         """
         lambdas = self._matrix_eigenvalues(matrix)
+        logger.debug("lambdas: %s", lambdas)
         lambdas = self._convert_lambda_units(lambdas)
+        logger.debug("lambdas converted units: %s", lambdas)
         if matrix_type == "force" and flexible > 0:
             lambdas = self._flexible_dihedral(lambdas, flexible)
+            logger.debug("lambdas flexible halved: %s", lambdas)
 
         freqs = self._frequencies_from_lambdas(lambdas, temp)
         if freqs.size == 0:

CodeEntropy/levels/dihedrals.py

@@ -82,7 +82,7 @@ def build_conformational_states(
         states_ua: Dict[UAKey, List[str]] = {}
         states_res: List[List[str]] = [None] * number_groups
         flexible_ua: Dict[UAKey, int] = {}
-        flexible_res: List[int] = [None] * number_groups
+        flexible_res: List[int] = [0] * number_groups
 
         task = None
         if progress is not None:

tests/unit/CodeEntropy/entropy/nodes/test_vibrational_node.py

@@ -22,6 +22,7 @@ def shared_data_base():
         "reduced_universe": reduced_universe,
         "force_covariances": {"ua": {}, "res": [], "poly": []},
         "torque_covariances": {"ua": {}, "res": [], "poly": []},
+        "flexible_dihedrals": {"ua": 0, "res": [0]},
         "n_frames": 5,
         "reporter": MagicMock(),
     }
@@ -42,6 +43,7 @@ def shared_groups():
         "reduced_universe": ru,
         "force_covariances": {"ua": {}, "res": [], "poly": []},
         "torque_covariances": {"ua": {}, "res": [], "poly": []},
+        "flexible_dihedrals": {"ua": [], "res": []},
         "n_frames": 5,
         "reporter": MagicMock(),
     }
@@ -131,6 +133,7 @@ def test_run_raises_on_unknown_level(shared_data, monkeypatch):
 
     shared_data["force_covariances"] = {"ua": {}, "res": [], "poly": []}
     shared_data["torque_covariances"] = {"ua": {}, "res": [], "poly": []}
+    shared_data["flexible_dihedrals"] = {"ua": [], "res": []}
 
     with pytest.raises(ValueError):
         node.run(shared_data)
@@ -143,6 +146,7 @@ def test_run_united_atom_branch_stores_results(shared_data, monkeypatch):
     shared_data["groups"] = {0: [0]}
     shared_data["force_covariances"] = {"ua": {}, "res": [], "poly": []}
     shared_data["torque_covariances"] = {"ua": {}, "res": [], "poly": []}
+    shared_data["flexible_dihedrals"] = {"ua": [], "res": []}
 
     fake_pair = MagicMock(trans=1.0, rot=2.0)
     monkeypatch.setattr(
@@ -164,6 +168,7 @@ def test_unknown_level_raises(shared_data):
     shared_data["groups"] = {0: [0]}
     shared_data["force_covariances"] = {"ua": {}, "res": [], "poly": []}
     shared_data["torque_covariances"] = {"ua": {}, "res": [], "poly": []}
+    shared_data["flexible_dihedrals"] = {"ua": [], "res": []}
 
     with pytest.raises(ValueError):
         node.run(shared_data)
@@ -177,6 +182,7 @@ def test_polymer_branch_executes(shared_data, monkeypatch):
 
     shared_data["force_covariances"] = {"ua": {}, "res": [], "poly": [MagicMock()]}
     shared_data["torque_covariances"] = {"ua": {}, "res": [], "poly": [MagicMock()]}
+    shared_data["flexible_dihedrals"] = {"ua": [], "res": []}
 
     shared_data["reduced_universe"].atoms.fragments = [MagicMock(residues=[])]
 
@@ -205,6 +211,7 @@ def test_run_skips_empty_mol_ids_group():
         "reduced_universe": MagicMock(atoms=MagicMock(fragments=[])),
         "force_covariances": {"ua": {}, "res": [], "poly": []},
         "torque_covariances": {"ua": {}, "res": [], "poly": []},
+        "flexible_dihedrals": {"ua": [], "res": []},
         "n_frames": 5,
         "reporter": None,
     }
@@ -235,6 +242,7 @@ def test_compute_united_atom_entropy_logs_residue_data_when_reporter_present():
         residues=residues,
         force_ua={},
         torque_ua={},
+        flexible_ua={(7, 0): 0, (7, 1): 0},
         ua_frame_counts={(7, 0): 3, (7, 1): 4},
         reporter=reporter,
         n_frames_default=10,
@@ -255,6 +263,7 @@ def test_compute_force_torque_entropy_success_calls_vibrational_engine():
         temp=298.0,
         fmat=np.eye(3),
         tmat=np.eye(3),
+        flexible=0,
         highest=False,
     )
 
@@ -267,7 +276,7 @@ def test_compute_ft_entropy_success_calls_vibrational_engine_for_trans_and_rot()
     ve = MagicMock()
     ve.vibrational_entropy_calculation.side_effect = [1.5, 2.5]
 
-    out = node._compute_ft_entropy(ve=ve, temp=298.0, ftmat=np.eye(6))
+    out = node._compute_ft_entropy(ve=ve, temp=298.0, ftmat=np.eye(6), flexible=0)
 
     assert out == EntropyPair(trans=1.5, rot=2.5)
     assert ve.vibrational_entropy_calculation.call_count == 2
@@ -312,7 +321,7 @@ def test_compute_force_torque_entropy_returns_zero_when_missing_matrix(shared_gr
     node = VibrationalEntropyNode()
     ve = MagicMock()
     pair = node._compute_force_torque_entropy(
-        ve=ve, temp=298.0, fmat=None, tmat=np.eye(3), highest=True
+        ve=ve, temp=298.0, fmat=None, tmat=np.eye(3), flexible=0, highest=True
     )
     assert pair == EntropyPair(trans=0.0, rot=0.0)
 
@@ -326,17 +335,17 @@ def test_compute_force_torque_entropy_returns_zero_when_filter_removes_all(monke
     )
 
     pair = node._compute_force_torque_entropy(
-        ve=ve, temp=298.0, fmat=np.eye(3), tmat=np.eye(3), highest=True
+        ve=ve, temp=298.0, fmat=np.eye(3), tmat=np.eye(3), flexible=0, highest=True
     )
     assert pair == EntropyPair(trans=0.0, rot=0.0)
 
 
 def test_compute_ft_entropy_returns_zero_when_none():
     node = VibrationalEntropyNode()
     ve = MagicMock()
-    assert node._compute_ft_entropy(ve=ve, temp=298.0, ftmat=None) == EntropyPair(
-        trans=0.0, rot=0.0
-    )
+    assert node._compute_ft_entropy(
+        ve=ve, temp=298.0, ftmat=None, flexible=0
+    ) == EntropyPair(trans=0.0, rot=0.0)
 
 
 def test_log_molecule_level_results_ft_labels_branch():
@@ -376,7 +385,7 @@ def test_compute_ft_entropy_returns_zeros_when_filtered_ft_matrix_is_empty(monke
         lambda _arr, atol: np.empty((0, 0), dtype=float),
     )
 
-    out = node._compute_ft_entropy(ve=ve, temp=298.0, ftmat=np.eye(6))
+    out = node._compute_ft_entropy(ve=ve, temp=298.0, ftmat=np.eye(6), flexible=0)
 
     assert out == EntropyPair(trans=0.0, rot=0.0)
     ve.vibrational_entropy_calculation.assert_not_called()

tests/unit/CodeEntropy/entropy/test_vibrational_entropy_math.py

@@ -47,7 +47,7 @@ def test_entropy_components_returns_empty_when_all_invalid(run_manager):
     ve = VibrationalEntropy(run_manager=run_manager)
 
     ve._matrix_eigenvalues = lambda m: np.array([-1.0, 0.0, 0.0])
-    comps = ve._entropy_components(np.eye(3), temp=298.0)
+    comps = ve._entropy_components(np.eye(3), temp=298.0, matrix_type="", flexible=0)
 
     assert comps.size == 0
 
@@ -154,7 +154,7 @@ def test_vibrational_entropy_calculation_end_to_end_returns_float(
     )
 
     out = ve.vibrational_entropy_calculation(
-        np.eye(3), matrix_type="torque", temp=298.0, highest_level=False
+        np.eye(3), matrix_type="torque", temp=298.0, highest_level=False, flexible=0
     )
 
     assert isinstance(out, float)

tests/unit/CodeEntropy/levels/nodes/test_conformations_node.py

@@ -9,7 +9,7 @@ def test_compute_conformational_states_node_runs_and_writes_shared_data():
     node = ComputeConformationalStatesNode(universe_operations=uops)
 
     node._dihedral_analysis.build_conformational_states = MagicMock(
-        return_value=({"ua_key": ["0", "1"]}, [["00", "01"]])
+        return_value=({"ua_key": ["0", "1"]}, [["00", "01"]], {"ua_key": [0]}, [0])
     )
 
     shared = {

tests/unit/CodeEntropy/levels/test_dihedrals.py

@@ -209,7 +209,7 @@ def run(self):
             return SimpleNamespace(results=SimpleNamespace(angles=angles))
 
     with patch("CodeEntropy.levels.dihedrals.Dihedral", _FakeDihedral):
-        states = dt._assign_states(
+        states, num_flexible = dt._assign_states(
             data_container=MagicMock(),
             molecules=[0, 1],
             dihedrals=["D0"],
@@ -227,8 +227,10 @@ def test_assign_states_for_group_sets_empty_lists_and_delegates_for_nonempty():
 
     states_ua = {}
     states_res = [None, None]
+    flexible_ua = {}
+    flexible_res = [None, None]
 
-    with patch.object(dt, "_assign_states", return_value=["x"]) as assign_spy:
+    with patch.object(dt, "_assign_states", return_value=[["x"], 0]) as assign_spy:
         dt._assign_states_for_group(
             data_container=MagicMock(),
             group_id=1,
@@ -243,6 +245,8 @@ def test_assign_states_for_group_sets_empty_lists_and_delegates_for_nonempty():
             level_list=["united_atom", "residue"],
             states_ua=states_ua,
             states_res=states_res,
+            flexible_ua=flexible_ua,
+            flexible_res=flexible_res,
         )
 
     assert states_ua[(1, 0)] == []
@@ -264,7 +268,7 @@ def test_build_conformational_states_runs_group_and_skips_empty_group(monkeypatc
     monkeypatch.setattr(dt, "_collect_peaks_for_group", lambda **kw: ([], []))
     monkeypatch.setattr(dt, "_assign_states_for_group", lambda **kw: None)
 
-    states_ua, states_res = dt.build_conformational_states(
+    states_ua, states_res, flex_ua, flex_res = dt.build_conformational_states(
         data_container=MagicMock(),
         levels=levels,
         groups=groups,
@@ -337,7 +341,7 @@ def run(self):
             return SimpleNamespace(results=SimpleNamespace(angles=[]))
 
     with patch("CodeEntropy.levels.dihedrals.Dihedral", _FakeDihedral):
-        out = dt._assign_states(
+        out_state, out_flex = dt._assign_states(
             data_container=MagicMock(),
             molecules=[0],
             dihedrals=[],
@@ -347,7 +351,7 @@ def run(self):
             step=1,
         )
 
-    assert out == []
+    assert out_state == []
 
 
 def test_identify_peaks_multiple_molecules_real_histogram():
@@ -420,8 +424,10 @@ def test_assign_states_for_group_residue_nonempty_calls_assign_states():
 
     states_ua = {}
     states_res = [None, None]
+    flexible_ua = {}
+    flexible_res = [None, None]
 
-    with patch.object(dt, "_assign_states", return_value=["A"]) as spy:
+    with patch.object(dt, "_assign_states", return_value=[["A"], 0]) as spy:
         dt._assign_states_for_group(
             data_container=MagicMock(),
             group_id=1,
@@ -436,6 +442,8 @@ def test_assign_states_for_group_residue_nonempty_calls_assign_states():
             level_list=["residue"],
             states_ua=states_ua,
             states_res=states_res,
+            flexible_ua=flexible_ua,
+            flexible_res=flexible_res,
         )
 
     assert states_res[1] == ["A"]
@@ -463,7 +471,7 @@ def run(self):
             return SimpleNamespace(results=SimpleNamespace(angles=angles))
 
     with patch("CodeEntropy.levels.dihedrals.Dihedral", _FakeDihedral):
-        states = dt._assign_states(
+        states, num_flex = dt._assign_states(
             data_container=MagicMock(),
             molecules=[0, 1],
             dihedrals=["D0"],
@@ -523,7 +531,7 @@ def run(self):
             return SimpleNamespace(results=SimpleNamespace(angles=angles))
 
     with patch("CodeEntropy.levels.dihedrals.Dihedral", _FakeDihedral):
-        states = dt._assign_states(
+        states, num_flex = dt._assign_states(
             data_container=MagicMock(),
             molecules=[0],
             dihedrals=["D0"],
@@ -571,7 +579,7 @@ def test_build_conformational_states_with_progress_skips_empty_molecule_group():
     groups = {0: []}
     levels = {}
 
-    states_ua, states_res = dt.build_conformational_states(
+    states_ua, states_res, flex_ua, flex_res = dt.build_conformational_states(
         data_container=MagicMock(),
         levels=levels,
         groups=groups,