docs: tidy comments and standardise to Google-style docstrings

Author: harryswift01

CodeEntropy/config/runtime.py

@@ -383,7 +383,6 @@ def read_universe(self, path: str) -> mda.Universe:
     def change_lambda_units(self, arg_lambdas: Any) -> Any:
         """Unit of lambdas : kJ2 mol-2 A-2 amu-1
         change units of lambda to J/s2"""
-        # return arg_lambdas * N_AVOGADRO * N_AVOGADRO * AMU2KG * 1e-26
         return arg_lambdas * 1e29 / self.N_AVOGADRO
 
     def get_KT2J(self, arg_temper: float) -> float:

CodeEntropy/entropy/configurational.py

@@ -96,7 +96,7 @@ def assign_conformation(
         Raises:
             ValueError: If `bin_width` or `step` are invalid.
         """
-        _ = number_frames  # kept for compatibility; sizing follows the slice length.
+        _ = number_frames
 
         config = ConformationConfig(
             bin_width=int(bin_width),
@@ -115,7 +115,6 @@ def assign_conformation(
         peak_values = self._find_histogram_peaks(phi, config.bin_width)
 
         if peak_values.size == 0:
-            # No peaks means no distinguishable states; assign everything to 0.
             return np.zeros(n_slice, dtype=int)
 
         states = self._assign_nearest_peaks(phi, peak_values)
@@ -138,21 +137,18 @@ def conformational_entropy_calculation(
         Returns:
             Conformational entropy in J/mol/K.
         """
-        _ = number_frames  # accepted as metadata; distribution uses observed counts.
+        _ = number_frames
 
         arr = self._to_1d_array(states)
         if arr is None or arr.size == 0:
             return 0.0
 
-        # If states contain only falsy values (e.g., all zeros) this is still valid:
-        # entropy would be 0 because only one state is present.
         values, counts = np.unique(arr, return_counts=True)
         total_count = int(np.sum(counts))
         if total_count <= 0 or values.size <= 1:
             return 0.0
 
         probs = counts.astype(float) / float(total_count)
-        # Guard against log(0) (shouldn't happen because counts>0), but keep robust.
         probs = probs[probs > 0.0]
 
         s_conf = -self._GAS_CONST * float(np.sum(probs * np.log(probs)))
@@ -174,7 +170,6 @@ def _validate_assignment_config(config: ConformationConfig) -> None:
         if config.bin_width <= 0 or config.bin_width > 360:
             raise ValueError("bin_width must be in the range (0, 360]")
         if 360 % config.bin_width != 0:
-            # Not strictly required, but prevents uneven bins and edge-case confusion.
             logger.warning(
                 "bin_width=%s does not evenly divide 360; histogram bins will be "
                 "uneven.",
@@ -242,8 +237,6 @@ def _assign_nearest_peaks(phi: np.ndarray, peak_values: np.ndarray) -> np.ndarra
         Returns:
             Integer state labels aligned with `phi`.
         """
-        # Vectorized nearest-peak assignment
-        # shape: (n_frames, n_peaks)
         distances = np.abs(phi[:, None] - peak_values[None, :])
         return np.argmin(distances, axis=1).astype(int)
 

CodeEntropy/entropy/nodes/vibrational.py

@@ -21,20 +21,53 @@
 
 @dataclass(frozen=True)
 class EntropyPair:
-    """Container for paired translational and rotational entropy values."""
+    """Container for paired translational and rotational entropy values.
+
+    Attributes:
+        trans: Translational vibrational entropy value.
+        rot: Rotational vibrational entropy value.
+    """
 
     trans: float
     rot: float
 
 
 class VibrationalEntropyNode:
-    """Compute vibrational entropy from force/torque (and optional FT) covariances."""
+    """Compute vibrational entropy from force/torque (and optional FT) covariances.
+
+    This node reads covariance matrices from a shared data mapping, computes
+    translational and rotational vibrational entropy at requested hierarchy levels,
+    and stores results back into the shared data structure.
+
+    The node supports:
+      - Force and torque covariance matrices ("force" / "torque") at residue/polymer
+        levels.
+      - United-atom per-residue covariances keyed by (group_id, residue_id).
+      - Optional combined force-torque covariance matrices ("forcetorque") for the
+        highest level when enabled via args.combined_forcetorque.
+    """
 
     def __init__(self) -> None:
+        """Initialize the node with matrix utilities and numerical tolerances."""
         self._mat_ops = MatrixUtils()
         self._zero_atol = 1e-8
 
     def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> Dict[str, Any]:
+        """Run vibrational entropy calculations and update the shared data mapping.
+
+        Args:
+            shared_data: Mutable mapping containing inputs (covariances, groups,
+                levels, args, etc.) and where outputs will be written.
+            **_: Unused keyword arguments, accepted for framework compatibility.
+
+        Returns:
+            A dict containing the computed vibrational entropy results under the
+            key "vibrational_entropy".
+
+        Raises:
+            ValueError: If an unknown level is encountered in the level list for a
+                representative molecule.
+        """
         ve = self._build_entropy_engine(shared_data)
         temp = shared_data["args"].temperature
 
@@ -125,18 +158,48 @@ def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> Dict[str, Any]
     def _build_entropy_engine(
         self, shared_data: Mapping[str, Any]
     ) -> VibrationalEntropy:
+        """Construct the vibrational entropy engine used for calculations.
+
+        Args:
+            shared_data: Read-only mapping containing a "run_manager" entry.
+
+        Returns:
+            A configured VibrationalEntropy instance.
+        """
         return VibrationalEntropy(
             run_manager=shared_data["run_manager"],
         )
 
     def _get_group_id_to_index(self, shared_data: Mapping[str, Any]) -> Dict[int, int]:
+        """Return a mapping from group_id to contiguous index used by covariance lists.
+
+        If a precomputed mapping is provided under "group_id_to_index", it is used.
+        Otherwise, the mapping is derived from the insertion order of "groups".
+
+        Args:
+            shared_data: Read-only mapping containing "groups" and optionally
+                "group_id_to_index".
+
+        Returns:
+            Dictionary mapping each group_id to an integer index.
+        """
         gid2i = shared_data.get("group_id_to_index")
         if isinstance(gid2i, dict) and gid2i:
             return gid2i
         groups = shared_data["groups"]
         return {gid: i for i, gid in enumerate(groups.keys())}
 
     def _get_ua_frame_counts(self, shared_data: Mapping[str, Any]) -> Dict[CovKey, int]:
+        """Extract per-(group,residue) frame counts for united-atom covariances.
+
+        Args:
+            shared_data: Read-only mapping which may contain nested frame count data
+                under shared_data["frame_counts"]["ua"].
+
+        Returns:
+            A dict keyed by (group_id, residue_id) containing frame counts. Returns
+            an empty dict if not present or not well-formed.
+        """
         counts = shared_data.get("frame_counts", {})
         if isinstance(counts, dict):
             ua_counts = counts.get("ua", {})
@@ -158,6 +221,27 @@ def _compute_united_atom_entropy(
         n_frames_default: int,
         highest: bool,
     ) -> EntropyPair:
+        """Compute total united-atom vibrational entropy for a group's residues.
+
+        Iterates over residues, looks up per-residue force and torque covariance
+        matrices keyed by (group_id, residue_index), computes entropy contributions,
+        accumulates totals, and optionally reports per-residue values.
+
+        Args:
+            ve: VibrationalEntropy calculation engine.
+            temp: Temperature (K) for entropy calculation.
+            group_id: Identifier for the group being processed.
+            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.
+            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.
+            highest: Whether this computation is at the highest requested level.
+
+        Returns:
+            EntropyPair with summed translational and rotational entropy across residues
+        """
         s_trans_total = 0.0
         s_rot_total = 0.0
 
@@ -207,6 +291,22 @@ def _compute_force_torque_entropy(
         tmat: Any,
         highest: bool,
     ) -> EntropyPair:
+        """Compute vibrational entropy from separate force and torque covariances.
+
+        Matrices are filtered to remove (near-)zero rows/columns before computation.
+        If either matrix is missing or becomes empty after filtering, returns zeros.
+
+        Args:
+            ve: VibrationalEntropy calculation engine.
+            temp: Temperature (K) for entropy calculation.
+            fmat: Force covariance matrix (array-like) or None.
+            tmat: Torque covariance matrix (array-like) or None.
+            highest: Whether this computation is at the highest requested level.
+
+        Returns:
+            EntropyPair containing translational entropy (from force covariance) and
+            rotational entropy (from torque covariance).
+        """
         if fmat is None or tmat is None:
             return EntropyPair(trans=0.0, rot=0.0)
 
@@ -235,6 +335,20 @@ def _compute_ft_entropy(
         temp: float,
         ftmat: Any,
     ) -> EntropyPair:
+        """Compute vibrational entropy from a combined force-torque covariance matrix.
+
+        The combined covariance matrix is filtered to remove (near-)zero rows/columns
+        before computation. If missing or empty after filtering, returns zeros.
+
+        Args:
+            ve: VibrationalEntropy calculation engine.
+            temp: Temperature (K) for entropy calculation.
+            ftmat: Combined force-torque covariance matrix (array-like) or None.
+
+        Returns:
+            EntropyPair containing translational and rotational entropy values derived
+            from the combined covariance matrix.
+        """
         if ftmat is None:
             return EntropyPair(trans=0.0, rot=0.0)
 
@@ -259,6 +373,14 @@ def _store_results(
         level: str,
         pair: EntropyPair,
     ) -> None:
+        """Store entropy results for a group/level into the results structure.
+
+        Args:
+            results: Nested results dict indexed by group_id then level.
+            group_id: Group identifier to store under.
+            level: Hierarchy level name (e.g., "united_atom", "residue", "polymer").
+            pair: EntropyPair containing translational and rotational values.
+        """
         results[group_id][level] = {"trans": pair.trans, "rot": pair.rot}
 
     @staticmethod
@@ -270,6 +392,15 @@ def _log_molecule_level_results(
         *,
         use_ft_labels: bool,
     ) -> None:
+        """Log molecule-level entropy results to the reporter, if available.
+
+        Args:
+            reporter: Optional reporter object supporting add_results_data calls.
+            group_id: Group identifier being reported.
+            level: Hierarchy level name being reported.
+            pair: EntropyPair containing translational and rotational values.
+            use_ft_labels: Whether to use FT-specific labels for the entropy types.
+        """
         if reporter is None:
             return
 
@@ -285,6 +416,15 @@ def _log_molecule_level_results(
 
     @staticmethod
     def _get_indexed_matrix(mats: Any, index: int) -> Any:
+        """Safely retrieve mats[index] if mats is indexable and index is in range.
+
+        Args:
+            mats: Indexable container of matrices (e.g., list/tuple) or other object.
+            index: Desired index.
+
+        Returns:
+            The matrix at the given index if available; otherwise None.
+        """
         try:
             return mats[index] if index < len(mats) else None
         except TypeError:

CodeEntropy/entropy/orientational.py

@@ -95,7 +95,6 @@ def calculate(self, neighbours: Mapping[str, int]) -> OrientationalEntropyResult
         total = 0.0
         for species, count in neighbours.items():
             if self._is_water(species):
-                # Water handling can be added later (e.g., via a strategy).
                 logger.debug(
                     "Skipping water species %s in orientational entropy.", species
                 )
@@ -141,7 +140,6 @@ def _entropy_contribution(self, neighbour_count: int) -> float:
             return 0.0
 
         omega = self._omega(neighbour_count)
-        # omega should always be > 0 when neighbour_count > 0, but guard anyway.
         if omega <= 0.0:
             return 0.0
 

CodeEntropy/entropy/vibrational.py

@@ -200,7 +200,6 @@ def _entropy_components_from_frequencies(
         kT = float(self._run_manager.get_KT2J(temp))
         exponent = (self._planck_const * frequencies) / kT
 
-        # Numerically stable enough for typical ranges; callers filter eigenvalues.
         exp_pos = np.exp(exponent)
         exp_neg = np.exp(-exponent)
 

CodeEntropy/entropy/workflow.py

@@ -114,7 +114,6 @@ def execute(self) -> None:
         if self._args.water_entropy and water_groups:
             self._compute_water_entropy(traj, water_groups)
         else:
-            # If water entropy isn't computed, include water in the remaining groups.
             nonwater_groups.update(water_groups)
 
         shared_data = self._build_shared_data(
@@ -297,8 +296,6 @@ def _compute_water_entropy(
         water_entropy = WaterEntropy(self._args)
 
         for group_id in water_groups.keys():
-            # WaterEntropy currently exposes a concrete API; keep this manager
-            # as an orchestrator and avoid duplicating internals here.
             water_entropy._calculate_water_entropy(
                 universe=self._universe,
                 start=traj.start,
@@ -307,7 +304,6 @@ def _compute_water_entropy(
                 group_id=group_id,
             )
 
-        # Exclude water from subsequent analysis when water entropy has been computed.
         self._args.selection_string = (
             f"{self._args.selection_string} and not water"
             if self._args.selection_string != "all"

CodeEntropy/levels/dihedrals.py

@@ -10,13 +10,6 @@
 
 import numpy as np
 from MDAnalysis.analysis.dihedrals import Dihedral
-from rich.progress import (
-    BarColumn,
-    Progress,
-    SpinnerColumn,
-    TextColumn,
-    TimeElapsedColumn,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -497,19 +490,3 @@ def _assign_states(
 
         logger.debug("States: %s", states)
         return states
-
-    @staticmethod
-    def _count_total_items(levels, groups) -> int:
-        """Count total progress items."""
-        return sum(len(levels[mol_id]) for mols in groups.values() for mol_id in mols)
-
-    @staticmethod
-    def _progress_bar(total_items: int) -> Progress:
-        """Create a Rich progress bar."""
-        return Progress(
-            SpinnerColumn(),
-            TextColumn("[bold blue]{task.fields[title]}", justify="right"),
-            BarColumn(),
-            TextColumn("[progress.percentage]{task.percentage:>3.1f}%"),
-            TimeElapsedColumn(),
-        )