diff --git a/CodeEntropy/config/arg_config_manager.py b/CodeEntropy/config/arg_config_manager.py index 799b950..7596c41 100644 --- a/CodeEntropy/config/arg_config_manager.py +++ b/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}." ) diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py index 75acfe5..2be1434 100644 --- a/CodeEntropy/entropy.py +++ b/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) diff --git a/CodeEntropy/levels.py b/CodeEntropy/levels.py index 62523e6..89c185f 100644 --- a/CodeEntropy/levels.py +++ b/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: diff --git a/tests/test_CodeEntropy/test_levels.py b/tests/test_CodeEntropy/test_levels.py index 0779fdc..08d83f8 100644 --- a/tests/test_CodeEntropy/test_levels.py +++ b/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