Align CodeEntropy with entropy methods used in 2019 paper

CodeEntropy/config/arg_config_manager.py

@@ -85,6 +85,18 @@
         "help": "How to group molecules for averaging",
         "default": "molecules",
     },
+    "combined_forcetorque": {
+        "type": bool,
+        "help": """Use combined force-torque matrix for residue
+            level vibrational entropies""",
+        "default": True,
+    },
+    "customised_axes": {
+        "type": bool,
+        "help": """Use bonded axes to rotate forces for UA
+            level vibrational entropies""",
+        "default": True,
+    },
 }
 
 

CodeEntropy/entropy.py

@@ -122,7 +122,7 @@ def execute(self):
         else:
             nonwater_groups.update(water_groups)
 
-        force_matrices, torque_matrices, frame_counts = (
+        force_matrices, torque_matrices, forcetorque_matrices, frame_counts = (
             self._level_manager.build_covariance_matrices(
                 self,
                 reduced_atom,
@@ -133,6 +133,8 @@ def execute(self):
                 step,
                 number_frames,
                 self._args.force_partitioning,
+                self._args.combined_forcetorque,
+                self._args.customised_axes,
             )
         )
 
@@ -155,6 +157,7 @@ def execute(self):
             nonwater_groups,
             force_matrices,
             torque_matrices,
+            forcetorque_matrices,
             states_ua,
             states_res,
             frame_counts,
@@ -230,6 +233,7 @@ def _compute_entropies(
         groups,
         force_matrices,
         torque_matrices,
+        forcetorque_matrices,
         states_ua,
         states_res,
         frame_counts,
@@ -309,6 +313,7 @@ def _compute_entropies(
                         f"Level: {level}",
                     )
                     highest = level == levels[groups[group_id][0]][-1]
+                    forcetorque_matrix = None
 
                     if level == "united_atom":
                         self._process_united_atom_entropy(
@@ -326,6 +331,8 @@ def _compute_entropies(
                         )
 
                     elif level == "residue":
+                        if highest:
+                            forcetorque_matrix = forcetorque_matrices["res"][group_id]
                         self._process_vibrational_entropy(
                             group_id,
                             mol,
@@ -334,6 +341,7 @@ def _compute_entropies(
                             level,
                             force_matrices["res"][group_id],
                             torque_matrices["res"][group_id],
+                            forcetorque_matrix,
                             highest,
                         )
 
@@ -347,6 +355,8 @@ def _compute_entropies(
                         )
 
                     elif level == "polymer":
+                        if highest:
+                            forcetorque_matrix = forcetorque_matrices["poly"][group_id]
                         self._process_vibrational_entropy(
                             group_id,
                             mol,
@@ -355,6 +365,7 @@ def _compute_entropies(
                             level,
                             force_matrices["poly"][group_id],
                             torque_matrices["poly"][group_id],
+                            forcetorque_matrix,
                             highest,
                         )
 
@@ -530,6 +541,7 @@ def _process_vibrational_entropy(
         level,
         force_matrix,
         torque_matrix,
+        forcetorque_matrix,
         highest,
     ):
         """
@@ -548,21 +560,43 @@ 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)
+        if forcetorque_matrix is not None:
+            forcetorque_matrix = self._level_manager.filter_zero_rows_columns(
+                forcetorque_matrix
+            )
 
-        torque_matrix = self._level_manager.filter_zero_rows_columns(torque_matrix)
+            S_FTtrans = ve.vibrational_entropy_calculation(
+                forcetorque_matrix, "forcetorqueTRANS", self._args.temperature, highest
+            )
+            S_FTrot = ve.vibrational_entropy_calculation(
+                forcetorque_matrix, "forcetorqueROT", self._args.temperature, highest
+            )
 
-        # Calculate the vibrational entropy
-        S_trans = ve.vibrational_entropy_calculation(
-            force_matrix, "force", self._args.temperature, highest
-        )
-        S_rot = ve.vibrational_entropy_calculation(
-            torque_matrix, "torque", self._args.temperature, highest
-        )
+            self._data_logger.add_results_data(
+                group_id, level, "FTmat-Transvibrational", S_FTtrans
+            )
+            self._data_logger.add_results_data(
+                group_id, level, "FTmat-Rovibrational", S_FTrot
+            )
 
-        # Print the vibrational entropy for the molecule group
-        self._data_logger.add_results_data(group_id, level, "Transvibrational", S_trans)
-        self._data_logger.add_results_data(group_id, level, "Rovibrational", S_rot)
+        else:
+            force_matrix = self._level_manager.filter_zero_rows_columns(force_matrix)
+
+            torque_matrix = self._level_manager.filter_zero_rows_columns(torque_matrix)
+
+            # Calculate the vibrational entropy
+            S_trans = ve.vibrational_entropy_calculation(
+                force_matrix, "force", self._args.temperature, highest
+            )
+            S_rot = ve.vibrational_entropy_calculation(
+                torque_matrix, "torque", self._args.temperature, highest
+            )
+
+            # Print the vibrational entropy for the molecule group
+            self._data_logger.add_results_data(
+                group_id, level, "Transvibrational", S_trans
+            )
+            self._data_logger.add_results_data(group_id, level, "Rovibrational", S_rot)
 
         residue_group = "_".join(
             sorted(set(res.resname for res in mol_container.residues))
@@ -899,6 +933,7 @@ def vibrational_entropy_calculation(self, matrix, matrix_type, temp, highest_lev
         """
         # N beads at a level => 3N x 3N covariance matrix => 3N eigenvalues
         # Get eigenvalues of the given matrix and change units to SI units
+        logger.debug(f"matrix_type: {matrix_type}")
         lambdas = la.eigvals(matrix)
         logger.debug(f"Eigenvalues (lambdas) before unit change: {lambdas}")
 
@@ -939,6 +974,11 @@ def vibrational_entropy_calculation(self, matrix, matrix_type, temp, highest_lev
             else:
                 S_vib_total = sum(S_components[6:])
 
+        elif matrix_type == "forcetorqueTRANS":  # three lowest are translations
+            S_vib_total = sum(S_components[:3])
+        elif matrix_type == "forcetorqueROT":  # three highest are rotations
+            S_vib_total = sum(S_components[3:])
+
         else:  # torque covariance matrix - we always take all values into account
             S_vib_total = sum(S_components)