Merge pull request #193 from CCPBioSim/192-force-partitioning

Force partitioning

Author: harryswift01

CodeEntropy/config/arg_config_manager.py

@@ -268,5 +268,5 @@ def _check_input_force_partitioning(self, args):
         if args.force_partitioning != default_value:
             logger.warning(
                 f"'force_partitioning' is set to {args.force_partitioning},"
-                " which differs from the default ({default_value})."
+                f" which differs from the default {default_value}."
             )

CodeEntropy/entropy.py

@@ -118,6 +118,7 @@ def execute(self):
                 end,
                 step,
                 number_frames,
+                self._args.force_partitioning,
             )
         )
 
@@ -548,6 +549,7 @@ def _process_vibrational_entropy(
         """
         # Find the relevant force and torque matrices and tidy them up
         # by removing rows and columns that are all zeros
+
         force_matrix = self._level_manager.filter_zero_rows_columns(force_matrix)
 
         torque_matrix = self._level_manager.filter_zero_rows_columns(torque_matrix)

CodeEntropy/levels.py

@@ -83,6 +83,7 @@ def get_matrices(
         highest_level,
         force_matrix,
         torque_matrix,
+        force_partitioning,
     ):
         """
         Compute and accumulate force/torque covariance matrices for a given level.
@@ -94,6 +95,8 @@ def get_matrices(
           highest_level (bool): Whether this is the top (largest bead size) level.
           force_matrix, torque_matrix (np.ndarray or None): Accumulated matrices to add
           to.
+          force_partitioning (float): Factor to adjust force contributions,
+          default is 0.5.
 
         Returns:
           force_matrix (np.ndarray): Accumulated force covariance matrix.
@@ -119,10 +122,14 @@ def get_matrices(
 
             # Sort out coordinates, forces, and torques for each atom in the bead
             weighted_forces[bead_index] = self.get_weighted_forces(
-                data_container, list_of_beads[bead_index], trans_axes, highest_level
+                data_container,
+                list_of_beads[bead_index],
+                trans_axes,
+                highest_level,
+                force_partitioning,
             )
             weighted_torques[bead_index] = self.get_weighted_torques(
-                data_container, list_of_beads[bead_index], rot_axes
+                data_container, list_of_beads[bead_index], rot_axes, force_partitioning
             )
 
         # Create covariance submatrices
@@ -571,7 +578,7 @@ def get_sphCoord_axes(self, arg_r):
         return spherical_basis
 
     def get_weighted_forces(
-        self, data_container, bead, trans_axes, highest_level, force_partitioning=0.5
+        self, data_container, bead, trans_axes, highest_level, force_partitioning
     ):
         """
         Function to calculate the mass weighted forces for a given bead.
@@ -620,9 +627,7 @@ def get_weighted_forces(
 
         return weighted_force
 
-    def get_weighted_torques(
-        self, data_container, bead, rot_axes, force_partitioning=0.5
-    ):
+    def get_weighted_torques(self, data_container, bead, rot_axes, force_partitioning):
         """
         Function to calculate the moment of inertia weighted torques for a given bead.
 
@@ -747,6 +752,7 @@ def build_covariance_matrices(
         end,
         step,
         number_frames,
+        force_partitioning,
     ):
         """
         Construct average force and torque covariance matrices for all molecules and
@@ -770,6 +776,9 @@ def build_covariance_matrices(
             Step size for frame iteration.
         number_frames : int
             Total number of frames to process.
+        force_partitioning : float
+            Factor to adjust force contributions, default is 0.5.
+
 
         Returns
         -------
@@ -855,6 +864,7 @@ def build_covariance_matrices(
                                 force_avg,
                                 torque_avg,
                                 frame_counts,
+                                force_partitioning,
                             )
 
                             progress.advance(task)
@@ -873,6 +883,7 @@ def update_force_torque_matrices(
         force_avg,
         torque_avg,
         frame_counts,
+        force_partitioning,
     ):
         """
         Update the running averages of force and torque covariance matrices
@@ -913,7 +924,8 @@ def update_force_torque_matrices(
         frame_counts : dict
             Dictionary holding the count of frames processed for each molecule/level
             combination.
-
+        force_partitioning : float
+         Factor to adjust force contributions, default is 0.5.
         Returns
         -------
         None
@@ -946,6 +958,7 @@ def update_force_torque_matrices(
                     highest,
                     None if key not in force_avg["ua"] else force_avg["ua"][key],
                     None if key not in torque_avg["ua"] else torque_avg["ua"][key],
+                    force_partitioning,
                 )
 
                 if key not in force_avg["ua"]:
@@ -979,6 +992,7 @@ def update_force_torque_matrices(
                     if torque_avg[key][group_id] is None
                     else torque_avg[key][group_id]
                 ),
+                force_partitioning,
             )
 
             if force_avg[key][group_id] is None:

tests/test_CodeEntropy/test_levels.py

@@ -92,6 +92,7 @@ def test_get_matrices(self):
             highest_level=True,
             force_matrix=None,
             torque_matrix=None,
+            force_partitioning=0.5,
         )
 
         # Assertions
@@ -139,6 +140,7 @@ def test_get_matrices_force_shape_mismatch(self):
                 highest_level=True,
                 force_matrix=bad_force_matrix,
                 torque_matrix=correct_torque_matrix,
+                force_partitioning=0.5,
             )
 
         self.assertIn("Inconsistent force matrix shape", str(context.exception))
@@ -174,6 +176,7 @@ def test_get_matrices_torque_shape_mismatch(self):
                 highest_level=True,
                 force_matrix=correct_force_matrix,
                 torque_matrix=bad_torque_matrix,
+                force_partitioning=0.5,
             )
 
         self.assertIn("Inconsistent torque matrix shape", str(context.exception))
@@ -207,6 +210,7 @@ def test_get_matrices_torque_consistency(self):
             highest_level=True,
             force_matrix=initial_force_matrix.copy(),
             torque_matrix=initial_torque_matrix.copy(),
+            force_partitioning=0.5,
         )
 
         force_matrix_2, torque_matrix_2 = level_manager.get_matrices(
@@ -216,6 +220,7 @@ def test_get_matrices_torque_consistency(self):
             highest_level=True,
             force_matrix=initial_force_matrix.copy(),
             torque_matrix=initial_torque_matrix.copy(),
+            force_partitioning=0.5,
         )
 
         # Check that repeated calls produce the same output
@@ -692,7 +697,7 @@ def test_get_weighted_force_with_partitioning(self):
         trans_axes = np.identity(3)
 
         result = level_manager.get_weighted_forces(
-            data_container, bead, trans_axes, highest_level=True
+            data_container, bead, trans_axes, highest_level=True, force_partitioning=0.5
         )
 
         expected = (0.5 * np.array([2.0, 0.0, 0.0])) / np.sqrt(4.0)
@@ -717,7 +722,11 @@ def test_get_weighted_force_without_partitioning(self):
         trans_axes = np.identity(3)
 
         result = level_manager.get_weighted_forces(
-            data_container, bead, trans_axes, highest_level=False
+            data_container,
+            bead,
+            trans_axes,
+            highest_level=False,
+            force_partitioning=0.5,
         )
 
         expected = np.array([3.0, 0.0, 0.0]) / np.sqrt(1.0)
@@ -743,7 +752,11 @@ def test_get_weighted_forces_zero_mass_raises_value_error(self):
 
         with self.assertRaises(ValueError):
             level_manager.get_weighted_forces(
-                data_container, bead, trans_axes, highest_level=True
+                data_container,
+                bead,
+                trans_axes,
+                highest_level=True,
+                force_partitioning=0.5,
             )
 
     def test_get_weighted_forces_negative_mass_raises_value_error(self):
@@ -766,7 +779,11 @@ def test_get_weighted_forces_negative_mass_raises_value_error(self):
 
         with self.assertRaises(ValueError):
             level_manager.get_weighted_forces(
-                data_container, bead, trans_axes, highest_level=True
+                data_container,
+                bead,
+                trans_axes,
+                highest_level=True,
+                force_partitioning=0.5,
             )
 
     def test_get_weighted_torques_weighted_torque_basic(self):
@@ -795,7 +812,11 @@ def test_get_weighted_torques_weighted_torque_basic(self):
         # Rotation axes (identity matrix)
         rot_axes = np.identity(3)
 
-        result = level_manager.get_weighted_torques(data_container, bead, rot_axes)
+        force_partitioning = 0.5
+
+        result = level_manager.get_weighted_torques(
+            data_container, bead, rot_axes, force_partitioning
+        )
 
         np.testing.assert_array_almost_equal(result, np.array([0.0, 0.0, 0.5]))
 
@@ -821,7 +842,11 @@ def test_get_weighted_torques_zero_torque_skips_division(self):
 
         rot_axes = np.identity(3)
 
-        result = level_manager.get_weighted_torques(data_container, bead, rot_axes)
+        force_partitioning = 0.5
+
+        result = level_manager.get_weighted_torques(
+            data_container, bead, rot_axes, force_partitioning
+        )
         np.testing.assert_array_almost_equal(result, np.zeros(3))
 
     def test_get_weighted_torques_zero_moi_raises(self):
@@ -852,8 +877,12 @@ def test_get_weighted_torques_zero_moi_raises(self):
 
         rot_axes = np.identity(3)
 
+        force_partitioning = 0.5
+
         with self.assertRaises(ZeroDivisionError):
-            level_manager.get_weighted_torques(data_container, bead, rot_axes)
+            level_manager.get_weighted_torques(
+                data_container, bead, rot_axes, force_partitioning
+            )
 
     def test_get_weighted_torques_negative_moi_raises(self):
         """
@@ -883,8 +912,12 @@ def test_get_weighted_torques_negative_moi_raises(self):
 
         rot_axes = np.identity(3)
 
+        foce_partitioning = 0.5
+
         with self.assertRaises(ValueError) as context:
-            level_manager.get_weighted_torques(data_container, bead, rot_axes)
+            level_manager.get_weighted_torques(
+                data_container, bead, rot_axes, foce_partitioning
+            )
 
         self.assertIn(
             "Negative value encountered for moment of inertia", str(context.exception)
@@ -995,6 +1028,7 @@ def test_build_covariance_matrices_atomic(self):
             end=2,
             step=1,
             number_frames=2,
+            force_partitioning=0.5,
         )
 
         # Assert returned matrices are dictionaries with correct keys
@@ -1057,6 +1091,7 @@ def test_update_force_torque_matrices_united_atom(self):
             force_avg=force_avg,
             torque_avg=torque_avg,
             frame_counts=frame_counts,
+            force_partitioning=0.5,
         )
 
         expected_keys = [(0, 0), (0, 1)]
@@ -1107,6 +1142,7 @@ def test_update_force_torque_matrices_united_atom_increment(self):
             force_avg=force_avg,
             torque_avg=torque_avg,
             frame_counts=frame_counts,
+            force_partitioning=0.5,
         )
 
         # Second call: update
@@ -1123,6 +1159,7 @@ def test_update_force_torque_matrices_united_atom_increment(self):
             force_avg=force_avg,
             torque_avg=torque_avg,
             frame_counts=frame_counts,
+            force_partitioning=0.5,
         )
 
         expected_force = f_mat_1 + (f_mat_2 - f_mat_1) / 2
@@ -1162,6 +1199,7 @@ def test_update_force_torque_matrices_residue(self):
             force_avg=force_avg,
             torque_avg=torque_avg,
             frame_counts=frame_counts,
+            force_partitioning=0.5,
         )
 
         np.testing.assert_array_equal(force_avg["res"][0], f_mat_mock)
@@ -1206,6 +1244,7 @@ def test_update_force_torque_matrices_incremental_average(self):
             force_avg=force_avg,
             torque_avg=torque_avg,
             frame_counts=frame_counts,
+            force_partitioning=0.5,
         )
 
         # Second update
@@ -1220,6 +1259,7 @@ def test_update_force_torque_matrices_incremental_average(self):
             force_avg=force_avg,
             torque_avg=torque_avg,
             frame_counts=frame_counts,
+            force_partitioning=0.5,
         )
 
         expected_force = f_mat_1 + (f_mat_2 - f_mat_1) / 2