build_covariance_matrices in EntropyManager now includes force partitioning argument

Author: ioanaapapa

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
+        print(highest)
         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: