separating the building of conformational states from the covariance matrices and entropy calculations.

Author: skfegan

CodeEntropy/entropy.py

@@ -47,11 +47,25 @@ def execute(self):
         """
         start, end, step = self._get_trajectory_bounds()
         number_frames = self._get_number_frames(start, end, step)
+        ve = VibrationalEntropy(
+            self._run_manager,
+            self._args,
+            self._universe,
+            self._data_logger,
+            self._level_manager,
+        )
+        ce = ConformationalEntropy(
+            self._run_manager,
+            self._args,
+            self._universe,
+            self._data_logger,
+            self._level_manager,
+        )
 
         self._handle_water_entropy(start, end, step)
         reduced_atom, number_molecules, levels = self._initialize_molecules()
 
-        force_matrices, torque_matrices, states_ua, states_res = (
+        force_matrices, torque_matrices = (
             self._level_manager.build_covariance_matrices(
                 self,
                 reduced_atom,
@@ -64,6 +78,21 @@ def execute(self):
             )
         )
 
+        states_ua, states_res = (
+            self._level_manager.build_conformational_states(
+                self,
+                reduced_atom,
+                number_molecules,
+                levels,
+                start,
+                end,
+                step,
+                number_frames,
+                self._args.bin_width,
+                ce,
+            )
+        )
+
         self._compute_entropies(
             reduced_atom,
             number_molecules,
@@ -76,6 +105,8 @@ def execute(self):
             start,
             end,
             step,
+            ve,
+            ce,
         )
 
         self._finalize_molecule_results()
@@ -136,6 +167,8 @@ def _compute_entropies(
         start,
         end,
         step,
+        ve,
+        ce,
     ):
         """
         Compute vibrational and conformational entropies for all molecules and levels.
@@ -163,51 +196,13 @@ def _compute_entropies(
             step (int): Step size for frame iteration.
         """
         bin_width = self._args.bin_width
-        ve = VibrationalEntropy(
-            self._run_manager,
-            self._args,
-            self._universe,
-            self._data_logger,
-            self._level_manager,
-        )
-        ce = ConformationalEntropy(
-            self._run_manager,
-            self._args,
-            self._universe,
-            self._data_logger,
-            self._level_manager,
-        )
 
         for mol_id in range(number_molecules):
             mol = self._get_molecule_container(reduced_atom, mol_id)
             for level in levels[mol_id]:
                 highest = level == levels[mol_id][-1]
 
                 if level == "united_atom":
-                    for res_id, residue in enumerate(mol.residues):
-                        key = (mol_id, res_id)
-                        res_container = self._run_manager.new_U_select_atom(
-                            mol,
-                            f"index {residue.atoms.indices[0]}:"
-                            f"{residue.atoms.indices[-1]}",
-                        )
-                        heavy_res = self._run_manager.new_U_select_atom(
-                            res_container, "not name H*"
-                        )
-
-                        states_ua[key] = (
-                            self._level_manager.compute_dihedral_conformations(
-                                heavy_res,
-                                level,
-                                ce,
-                                number_frames,
-                                bin_width,
-                                start,
-                                end,
-                                step,
-                            )
-                        )
-
                     self._process_united_atom_entropy(
                         mol_id,
                         mol,
@@ -222,10 +217,6 @@ def _compute_entropies(
                     )
 
                 elif level == "residue":
-                    states_res = self._level_manager.compute_dihedral_conformations(
-                        mol, level, ce, number_frames, bin_width, start, end, step
-                    )
-
                     self._process_vibrational_entropy(
                         mol_id,
                         mol,

CodeEntropy/levels.py

@@ -46,7 +46,7 @@ def select_levels(self, data_container):
 
         # fragments is MDAnalysis terminology for what chemists would call molecules
         number_molecules = len(data_container.atoms.fragments)
-        logger.debug("The number of molecules is {}.".format(number_molecules))
+        logger.debug(f"The number of molecules is {number_molecules}.")
 
         fragments = data_container.atoms.fragments
         levels = [[] for _ in range(number_molecules)]
@@ -265,20 +265,19 @@ def compute_dihedral_conformations(
         self,
         selector,
         level,
-        ce,
         number_frames,
         bin_width,
         start,
         end,
         step,
+        ce,
     ):
         """
         Compute dihedral conformations for a given selector and entropy level.
 
         Parameters:
             selector (AtomGroup): Atom selection to compute dihedrals for.
             level (str): Entropy level ("united_atom" or "residue").
-            ce (ConformationalEntropy): Conformational entropy calculator.
             number_frames (int): Number of frames to process.
             bin_width (float): Bin width for dihedral angle discretization.
             start (int): Start frame index.
@@ -750,8 +749,6 @@ def build_covariance_matrices(
             tuple: A tuple containing:
                 - force_matrices (dict): Force covariance matrices by level.
                 - torque_matrices (dict): Torque covariance matrices by level.
-                - states_ua (dict): Conformational states at the united-atom level.
-                - states_res (list): Conformational states at the residue level.
         """
         force_matrices = {
             "ua": {},
@@ -763,8 +760,6 @@ def build_covariance_matrices(
             "res": [None] * number_molecules,
             "poly": [None] * number_molecules,
         }
-        states_ua = {}
-        states_res = [None] * number_molecules
 
         for timestep in reduced_atom.trajectory[start:end:step]:
             time_index = timestep.frame - start
@@ -784,7 +779,7 @@ def build_covariance_matrices(
                         torque_matrices,
                     )
 
-        return force_matrices, torque_matrices, states_ua, states_res
+        return force_matrices, torque_matrices
 
     def update_force_torque_matrices(
         self,
@@ -896,3 +891,79 @@ def filter_zero_rows_columns(self, arg_matrix):
         logger.debug(f"arg_matrix: {arg_matrix}")
 
         return arg_matrix
+
+    def build_conformational_states(
+        self,
+        entropy_manager,
+        reduced_atom,
+        number_molecules,
+        levels,
+        start,
+        end,
+        step,
+        number_frames,
+        bin_width,
+        ce,
+    ):
+        """
+        Construct the conformational states for each molecule at
+        relevant levels.
+
+        Parameters:
+            entropy_manager (EntropyManager): Instance of the EntropyManager
+            reduced_atom (Universe): The reduced atom selection.
+            number_molecules (int): Number of molecules in the system.
+            levels (list): List of entropy levels per molecule.
+            start (int): Start frame index.
+            end (int): End frame index.
+            step (int): Step size for frame iteration.
+            number_frames (int): Total number of frames to process.
+
+        Returns:
+            tuple: A tuple containing:
+                - states_ua (dict): Conformational states at the united-atom level.
+                - states_res (list): Conformational states at the residue level.
+        """
+        states_ua = {}
+        states_res = [None] * number_molecules
+
+        for mol_id in range(number_molecules):
+            mol = entropy_manager._get_molecule_container(reduced_atom, mol_id)
+            for level in levels[mol_id]:
+                if level == "united_atom":
+                    for res_id, residue in enumerate(mol.residues):
+                        key = (mol_id, res_id)
+
+                        res_container = entropy_manager._run_manager.new_U_select_atom(
+                            mol,
+                            f"index {residue.atoms.indices[0]}:"
+                            f"{residue.atoms.indices[-1]}",
+                        )
+                        heavy_res = entropy_manager._run_manager.new_U_select_atom(
+                            res_container, "not name H*"
+                        )
+
+                        states_ua[key] = self.compute_dihedral_conformations(
+                            heavy_res,
+                            level,
+                            number_frames,
+                            bin_width,
+                            start,
+                            end,
+                            step,
+                            ce,
+                        )
+
+                if level == "res":
+                    states_res = self.compute_dihedral_conformations(
+                        mol,
+                        level,
+                        number_frames,
+                        bin_width,
+                        start,
+                        end,
+                        step,
+                        ce,
+                    )
+
+        return states_ua, states_res