Skip to content

Force partitioning #193

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
harryswift01 merged 4 commits into from
Nov 18, 2025
Merged

Force partitioning #193

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
37bacff
build_covariance_matrices in EntropyManager now includes force partit…
ioanaapapa Nov 18, 2025
a09b746
remove unecessary print statement used while debugging
ioanaapapa Nov 18, 2025
dde3684
warning now displays default force_partitioning value
ioanaapapa Nov 18, 2025
04499b9
fix tests in relation to #193 PR:
harryswift01 Nov 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion CodeEntropy/config/arg_config_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}."
)
2 changes: 2 additions & 0 deletions CodeEntropy/entropy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,7 @@ def execute(self):
end,
step,
number_frames,
self._args.force_partitioning,
)
)

Expand Down Expand Up @@ -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)
Expand Down
28 changes: 21 additions & 7 deletions CodeEntropy/levels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand All @@ -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
-------
Expand Down Expand Up @@ -855,6 +864,7 @@ def build_covariance_matrices(
force_avg,
torque_avg,
frame_counts,
force_partitioning,
)

progress.advance(task)
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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"]:
Expand Down Expand Up @@ -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:
Expand Down
56 changes: 48 additions & 8 deletions tests/test_CodeEntropy/test_levels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ def test_get_matrices(self):
highest_level=True,
force_matrix=None,
torque_matrix=None,
force_partitioning=0.5,
)

# Assertions
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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):
Expand All @@ -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):
Expand Down Expand Up @@ -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]))

Expand All @@ -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):
Expand Down Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down