Merge pull request #300 from CCPBioSim/278-flexible-dihedrals

Force Halving for Flexible Dihedrals

Author: skfegan

CodeEntropy/entropy/nodes/vibrational.py

@@ -75,6 +75,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)
 
@@ -109,6 +110,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_ua=flexible["ua"],
                         ua_frame_counts=ua_frame_counts,
                         reporter=reporter,
                         n_frames_default=shared_data.get("n_frames", 0),
@@ -123,11 +125,18 @@ def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> dict[str, Any]
                 if level in ("residue", "polymer"):
                     gi = gid2i[group_id]
 
+                    if level == "residue":
+                        flexible_res = flexible["res"][group_id]
+                    else:
+                        flexible_res = 0  # No polymer level flexible dihedrals
+
                     if combined and highest and ft_cov is not None:
                         ft_key = "res" if level == "residue" else "poly"
                         ftmat = self._get_indexed_matrix(ft_cov.get(ft_key, []), gi)
 
-                        pair = self._compute_ft_entropy(ve=ve, temp=temp, ftmat=ftmat)
+                        pair = self._compute_ft_entropy(
+                            ve=ve, temp=temp, ftmat=ftmat, flexible=flexible_res
+                        )
                         self._store_results(results, group_id, level, pair)
                         self._log_molecule_level_results(
                             reporter, group_id, level, pair, use_ft_labels=True
@@ -143,6 +152,7 @@ def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> dict[str, Any]
                         temp=temp,
                         fmat=fmat,
                         tmat=tmat,
+                        flexible=flexible_res,
                         highest=highest,
                     )
                     self._store_results(results, group_id, level, pair)
@@ -217,6 +227,7 @@ def _compute_united_atom_entropy(
         residues: Any,
         force_ua: Mapping[CovKey, Any],
         torque_ua: Mapping[CovKey, Any],
+        flexible_ua: Any,
         ua_frame_counts: Mapping[CovKey, int],
         reporter: Any | None,
         n_frames_default: int,
@@ -235,6 +246,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.
@@ -250,12 +262,14 @@ def _compute_united_atom_entropy(
             key = (group_id, res_id)
             fmat = force_ua.get(key)
             tmat = torque_ua.get(key)
+            flexible = flexible_ua.get(key)
 
             pair = self._compute_force_torque_entropy(
                 ve=ve,
                 temp=temp,
                 fmat=fmat,
                 tmat=tmat,
+                flexible=flexible,
                 highest=highest,
             )
 
@@ -290,6 +304,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.
@@ -322,10 +337,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))
 
@@ -335,6 +350,7 @@ def _compute_ft_entropy(
         ve: VibrationalEntropy,
         temp: float,
         ftmat: Any,
+        flexible: int,
     ) -> EntropyPair:
         """Compute vibrational entropy from a combined force-torque covariance matrix.
 
@@ -360,10 +376,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:
@@ -116,7 +123,12 @@ def _entropy_components(self, matrix: np.ndarray, temp: float) -> np.ndarray:
             Array of entropy components (J/mol/K) for each valid mode.
         """
         lambdas = self._matrix_eigenvalues(matrix)
+        logger.debug("lambdas: %s", lambdas)
         lambdas = self._convert_lambda_units(lambdas)
+        logger.debug("lambdas converted units: %s", lambdas)
+        if matrix_type == "force" and flexible > 0:
+            lambdas = self._flexible_dihedral(lambdas, flexible)
+            logger.debug("lambdas flexible halved: %s", lambdas)
 
         freqs = self._frequencies_from_lambdas(lambdas, temp)
         if freqs.size == 0:
@@ -149,6 +161,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

@@ -44,7 +44,7 @@ def build_conformational_states(
         step: int,
         bin_width: float,
         progress: _RichProgressSink | None = None,
-    ) -> tuple[dict[UAKey, list[str]], list[list[str]]]:
+    ) -> tuple[dict[UAKey, list[str]], list[list[str]], dict[UAKey, int], list[int]]:
         """Build conformational state labels from trajectory dihedrals.
 
         This method constructs discrete conformational state descriptors used in
@@ -72,7 +72,7 @@ def build_conformational_states(
                 and advance().
 
         Returns:
-            tuple: (states_ua, states_res)
+            tuple: (states_ua, states_res, flexible_ua, flexible_res)
 
             - states_ua: Dict mapping (group_id, local_residue_id) -> list of state
               labels (strings) across the analyzed trajectory.
@@ -87,6 +87,8 @@ def build_conformational_states(
         number_groups = len(groups)
         states_ua: dict[UAKey, list[str]] = {}
         states_res: list[list[str]] = [[] for _ in range(number_groups)]
+        flexible_ua: dict[UAKey, int] = {}
+        flexible_res: list[int] = [0] * number_groups
 
         task: TaskID | None = None
         if progress is not None:
@@ -149,12 +151,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 progress is not None and 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: Any, level_list: list[str]
@@ -268,6 +272,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(
@@ -282,6 +287,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(
@@ -353,7 +359,9 @@ def _identify_peaks(
             peaks = self._find_histogram_peaks(popul=popul, bin_value=bin_value)
             peak_values.append(peaks)
 
-            logger.debug(f"Dihedral: {dihedral_index}, Peak Values: {peak_values}")
+            logger.debug(
+                f"Dihedral: {dihedral_index}Peak Values: {peak_values[dihedral_index]}"
+            )
 
         return peak_values
 
@@ -404,6 +412,8 @@ def _assign_states_for_group(
         level_list: list[str],
         states_ua: dict[UAKey, list[str]],
         states_res: list[list[str]],
+        flexible_ua: dict[UAKey, list[int]],
+        flexible_res: list[int],
     ) -> None:
         """Assign UA and residue states for a group into output containers."""
         for level in level_list:
@@ -412,8 +422,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],
@@ -426,8 +437,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,
@@ -467,6 +479,7 @@ def _assign_states(
             List of state labels (strings).
         """
         states: list[str] = []
+        num_flexible = 0
 
         for molecule in molecules:
             conformations: list[list[Any]] = []
@@ -477,16 +490,29 @@ def _assign_states(
 
             for dihedral_index in range(len(dihedrals)):
                 conformation: list[Any] = []
+
+                # 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 (
@@ -501,4 +527,6 @@ def _assign_states(
             states.extend(mol_states)
 
         logger.debug(f"States: {states}")
-        return states
+        logger.debug(f"Number of flexible dihedrals: {num_flexible}")
+
+        return states, num_flexible

CodeEntropy/levels/neighbors.py

@@ -95,7 +95,8 @@ def get_neighbors(self, universe, levels, groups, search_type):
                 len(molecules) * number_frames
             )
             logger.debug(
-                "group: {group_id}number neighbors {average_number_neighbors[group_id]}"
+                f"group: {group_id}"
+                f"number neighbors {average_number_neighbors[group_id]}"
             )
 
         return average_number_neighbors
@@ -233,14 +234,13 @@ def _get_linear(self, rdkit_mol, number_heavy):
         Returns:
             linear (bool): True if molecule linear
         """
-        rdkit_heavy = Chem.RemoveHs(rdkit_mol)
-
         linear = False
         if number_heavy == 1:
             linear = False
         elif number_heavy == 2:
             linear = True
         else:
+            rdkit_heavy = Chem.RemoveHs(rdkit_mol)
             sp_count = 0
             for x in rdkit_heavy.GetAtoms():
                 if x.GetHybridization() == Chem.HybridizationType.SP: