adding flexible dihedrals to vibrational entropy calculation

Author: skfegan

CodeEntropy/entropy/nodes/vibrational.py

@@ -74,6 +74,7 @@ def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> Dict[str, Any]
         groups = shared_data["groups"]
         levels = shared_data["levels"]
         fragments = shared_data["reduced_universe"].atoms.fragments
+        flexible = shared_data["flexible_dihedrals"]
 
         gid2i = self._get_group_id_to_index(shared_data)
 
@@ -108,6 +109,7 @@ def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> Dict[str, Any]
                         residues=rep_mol.residues,
                         force_ua=force_cov["ua"],
                         torque_ua=torque_cov["ua"],
+                        flexible=flexible["ua"],
                         ua_frame_counts=ua_frame_counts,
                         reporter=reporter,
                         n_frames_default=shared_data.get("n_frames", 0),
@@ -137,11 +139,17 @@ def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> Dict[str, Any]
                     fmat = self._get_indexed_matrix(force_cov.get(cov_key, []), gi)
                     tmat = self._get_indexed_matrix(torque_cov.get(cov_key, []), gi)
 
+                    if level == "residue":
+                        flexible = flexible["res"][group_id]
+                    else:
+                        flexible = 0  # No polymer level flexible dihedrals
+
                     pair = self._compute_force_torque_entropy(
                         ve=ve,
                         temp=temp,
                         fmat=fmat,
                         tmat=tmat,
+                        flexible=flexible,
                         highest=highest,
                     )
                     self._store_results(results, group_id, level, pair)
@@ -216,6 +224,7 @@ def _compute_united_atom_entropy(
         residues: Any,
         force_ua: Mapping[CovKey, Any],
         torque_ua: Mapping[CovKey, Any],
+        flexible: Any,
         ua_frame_counts: Mapping[CovKey, int],
         reporter: Optional[Any],
         n_frames_default: int,
@@ -234,6 +243,7 @@ def _compute_united_atom_entropy(
             residues: Residue container/sequence for the representative molecule.
             force_ua: Mapping from (group_id, residue_id) to force covariance matrix.
             torque_ua: Mapping from (group_id, residue_id) to torque covariance matrix.
+            flexible: Data about number of flexible dihedrals
             ua_frame_counts: Mapping from (group_id, residue_id) to frame counts.
             reporter: Optional reporter object supporting add_residue_data calls.
             n_frames_default: Fallback frame count if per-residue count missing.
@@ -249,12 +259,14 @@ def _compute_united_atom_entropy(
             key = (group_id, res_id)
             fmat = force_ua.get(key)
             tmat = torque_ua.get(key)
+            flexible = flexible[key]
 
             pair = self._compute_force_torque_entropy(
                 ve=ve,
                 temp=temp,
                 fmat=fmat,
                 tmat=tmat,
+                flexible=flexible,
                 highest=highest,
             )
 
@@ -289,6 +301,7 @@ def _compute_force_torque_entropy(
         temp: float,
         fmat: Any,
         tmat: Any,
+        flexible: int,
         highest: bool,
     ) -> EntropyPair:
         """Compute vibrational entropy from separate force and torque covariances.
@@ -321,10 +334,10 @@ def _compute_force_torque_entropy(
             return EntropyPair(trans=0.0, rot=0.0)
 
         s_trans = ve.vibrational_entropy_calculation(
-            f, "force", temp, highest_level=highest
+            f, "force", temp, highest_level=highest, flexible=flexible
         )
         s_rot = ve.vibrational_entropy_calculation(
-            t, "torque", temp, highest_level=highest
+            t, "torque", temp, highest_level=highest, flexible=flexible
         )
         return EntropyPair(trans=float(s_trans), rot=float(s_rot))
 
@@ -334,6 +347,7 @@ def _compute_ft_entropy(
         ve: VibrationalEntropy,
         temp: float,
         ftmat: Any,
+        flexible: int,
     ) -> EntropyPair:
         """Compute vibrational entropy from a combined force-torque covariance matrix.
 
@@ -359,10 +373,10 @@ def _compute_ft_entropy(
             return EntropyPair(trans=0.0, rot=0.0)
 
         s_trans = ve.vibrational_entropy_calculation(
-            ft, "forcetorqueTRANS", temp, highest_level=True
+            ft, "forcetorqueTRANS", temp, highest_level=True, flexible=flexible
         )
         s_rot = ve.vibrational_entropy_calculation(
-            ft, "forcetorqueROT", temp, highest_level=True
+            ft, "forcetorqueROT", temp, highest_level=True, flexible=flexible
         )
         return EntropyPair(trans=float(s_trans), rot=float(s_rot))
 

CodeEntropy/entropy/vibrational.py

@@ -68,6 +68,7 @@ def vibrational_entropy_calculation(
         matrix_type: MatrixType,
         temp: float,
         highest_level: bool,
+        flexible: int,
     ) -> float:
         """Compute vibrational entropy for the given covariance matrix.
 
@@ -101,11 +102,17 @@ def vibrational_entropy_calculation(
         Raises:
             ValueError: If matrix_type is unknown.
         """
-        components = self._entropy_components(matrix, temp)
+        components = self._entropy_components(matrix, matrix_type, flexible, temp)
         total = self._sum_components(components, matrix_type, highest_level)
         return float(total)
 
-    def _entropy_components(self, matrix: np.ndarray, temp: float) -> np.ndarray:
+    def _entropy_components(
+        self,
+        matrix: np.ndarray,
+        matrix_type: str,
+        flexible: int,
+        temp: float,
+    ) -> np.ndarray:
         """Compute per-mode entropy components from a covariance matrix.
 
         Args:
@@ -117,6 +124,8 @@ def _entropy_components(self, matrix: np.ndarray, temp: float) -> np.ndarray:
         """
         lambdas = self._matrix_eigenvalues(matrix)
         lambdas = self._convert_lambda_units(lambdas)
+        if matrix_type == "force" and flexible > 0:
+            lambdas = self._flexible_dihedral(lambdas, flexible)
 
         freqs = self._frequencies_from_lambdas(lambdas, temp)
         if freqs.size == 0:
@@ -149,6 +158,26 @@ def _convert_lambda_units(self, lambdas: np.ndarray) -> np.ndarray:
         """
         return self._run_manager.change_lambda_units(lambdas)
 
+    def _flexible_dihedral(self, lambdas: np.ndarray, flexible: int) -> np.ndarray:
+        """Force halving for flexible dihedrals.
+
+        If N flexible dihedrals, halve the forces for the N largest eigenvalues.
+        The matrix has force^2 so use factor of 0.25 for eigenvalues.
+
+        Args:
+            lambdas: Eigenvalues
+            flexible: the number of flexible dihedrals in the molecule
+
+        Returns:
+            reduced lambdas
+        """
+        halved = sorted(lambdas, reverse=True)
+        for i in range(flexible):
+            halved[i] = 0.25 * halved[i]
+        lambdas = halved
+
+        return lambdas
+
     def _frequencies_from_lambdas(self, lambdas: np.ndarray, temp: float) -> np.ndarray:
         """Convert eigenvalues to frequencies with robust filtering.
 

CodeEntropy/levels/dihedrals.py

@@ -81,6 +81,8 @@ def build_conformational_states(
         number_groups = len(groups)
         states_ua: Dict[UAKey, List[str]] = {}
         states_res: List[List[str]] = [None] * number_groups
+        flexible_ua: Dict[UAKey, int] = {}
+        flexible_res: List[int] = [None] * number_groups
 
         task = None
         if progress is not None:
@@ -143,12 +145,14 @@ def build_conformational_states(
                 level_list=levels[molecules[0]],
                 states_ua=states_ua,
                 states_res=states_res,
+                flexible_ua=flexible_ua,
+                flexible_res=flexible_res,
             )
 
             if task is not None:
                 progress.advance(task)
 
-        return states_ua, states_res
+        return states_ua, states_res, flexible_ua, flexible_res
 
     def _collect_dihedrals_for_group(self, mol, level_list):
         """Collect UA and residue dihedral AtomGroups for a group.
@@ -260,6 +264,7 @@ def _collect_peaks_for_group(
             if level == "united_atom":
                 for res_id in range(len(dihedrals_ua)):
                     if len(dihedrals_ua[res_id]) == 0:
+                        # No dihedrals means no peaks
                         peaks_ua[res_id] = []
                     else:
                         peaks_ua[res_id] = self._identify_peaks(
@@ -274,6 +279,7 @@ def _collect_peaks_for_group(
 
             elif level == "residue":
                 if len(dihedrals_res) == 0:
+                    # No dihedrals means no peaks
                     peaks_res = []
                 else:
                     peaks_res = self._identify_peaks(
@@ -362,13 +368,13 @@ def _find_histogram_peaks(popul, bin_value):
             if bin_index == number_bins - 1:
                 if (
                     popul[bin_index] >= popul[bin_index - 1]
-                    and popul[bin_index] >= popul[0]
+                    and popul[bin_index] > popul[0]
                 ):
                     peaks.append(bin_value[bin_index])
             else:
                 if (
                     popul[bin_index] >= popul[bin_index - 1]
-                    and popul[bin_index] >= popul[bin_index + 1]
+                    and popul[bin_index] > popul[bin_index + 1]
                 ):
                     peaks.append(bin_value[bin_index])
 
@@ -389,6 +395,8 @@ def _assign_states_for_group(
         level_list,
         states_ua,
         states_res,
+        flexible_ua,
+        flexible_res,
     ):
         """Assign UA and residue states for a group into output containers."""
         for level in level_list:
@@ -397,8 +405,9 @@ def _assign_states_for_group(
                     key = (group_id, res_id)
                     if len(dihedrals_ua[res_id]) == 0:
                         states_ua[key] = []
+                        flexible_ua[key] = 0
                     else:
-                        states_ua[key] = self._assign_states(
+                        states_ua[key], flexible_ua[key] = self._assign_states(
                             data_container=data_container,
                             molecules=molecules,
                             dihedrals=dihedrals_ua[res_id],
@@ -411,8 +420,9 @@ def _assign_states_for_group(
             elif level == "residue":
                 if len(dihedrals_res) == 0:
                     states_res[group_id] = []
+                    flexible_res[group_id] = 0
                 else:
-                    states_res[group_id] = self._assign_states(
+                    states_res[group_id], flexible_res[group_id] = self._assign_states(
                         data_container=data_container,
                         molecules=molecules,
                         dihedrals=dihedrals_res,
@@ -452,6 +462,7 @@ def _assign_states(
             List of state labels (strings).
         """
         states = None
+        num_flexible = 0
 
         for molecule in molecules:
             conformations = []
@@ -462,16 +473,29 @@ def _assign_states(
 
             for dihedral_index in range(len(dihedrals)):
                 conformation = []
+
+                # Check for flexible dihedrals
+                if len(peaks[dihedral_index]) > 1 and molecule == 0:
+                    num_flexible += 1
+
+                # Get conformations
                 for timestep in range(number_frames):
                     value = dihedral_results.results.angles[timestep][dihedral_index]
+                    # We want postive values in range 0 to 360 to make
+                    # the peak assignment.
+                    # works using the fact that dihedrals have circular symmetry
+                    # (i.e. -15 degrees = +345 degrees)
                     if value < 0:
                         value += 360
 
+                    # Find the peak closest to the dihedral value
                     distances = [abs(value - peak) for peak in peaks[dihedral_index]]
                     conformation.append(np.argmin(distances))
 
                 conformations.append(conformation)
 
+            # Concatenate all the dihedrals in the molecule into the state
+            # for the frame.
             mol_states = [
                 state
                 for state in (
@@ -489,4 +513,5 @@ def _assign_states(
                 states.extend(mol_states)
 
         logger.debug("States: %s", states)
-        return states
+        logger.debug("Number of flexible dihedrals: %s", num_flexible)
+        return states, num_flexible

CodeEntropy/levels/nodes/conformations.py

@@ -15,6 +15,7 @@
 
 SharedData = Dict[str, Any]
 ConformationalStates = Dict[str, Any]
+FlexibleStates = Dict[str, Any]
 
 
 @dataclass(frozen=True)
@@ -39,10 +40,13 @@ class ComputeConformationalStatesNode:
 
     Produces:
         shared_data["conformational_states"] = {"ua": states_ua, "res": states_res}
+        shared_data["flexible_dihedrals"] = {"ua: flexible_ua, "res": flexible_res}
 
     Where:
         - states_ua is a dict keyed by (group_id, local_residue_id)
         - states_res is a list-like structure indexed by group_id (or equivalent)
+        - flexible_ua is a dict keyed by (group_id, local_residue_id)
+        - flexible_res is a list-like structure indexed by group_id (or equivalent)
     """
 
     def __init__(self, universe_operations: Any) -> None:
@@ -79,22 +83,33 @@ def run(
 
         cfg = self._build_config(shared_data)
 
-        states_ua, states_res = self._dihedral_analysis.build_conformational_states(
-            data_container=u,
-            levels=levels,
-            groups=groups,
-            start=cfg.start,
-            end=cfg.end,
-            step=cfg.step,
-            bin_width=cfg.bin_width,
-            progress=progress,
+        states_ua, states_res, flexible_ua, flexible_res = (
+            self._dihedral_analysis.build_conformational_states(
+                data_container=u,
+                levels=levels,
+                groups=groups,
+                start=cfg.start,
+                end=cfg.end,
+                step=cfg.step,
+                bin_width=cfg.bin_width,
+                progress=progress,
+            )
         )
 
+        # Get state information into shared_data
         conformational_states: ConformationalStates = {
             "ua": states_ua,
             "res": states_res,
         }
         shared_data["conformational_states"] = conformational_states
+
+        # Get flexible_dihedral data into shared_data
+        flexible_states: FlexibleStates = {
+            "ua": flexible_ua,
+            "res": flexible_res,
+        }
+        shared_data["flexible_dihedrals"] = flexible_states
+
         return {"conformational_states": conformational_states}
 
     @staticmethod