Introduce Basic Orientational Entropy Calculations

CodeEntropy/config/argparse.py

@@ -25,7 +25,7 @@
 import logging
 import os
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Set
+from typing import Any
 
 import yaml
 
@@ -48,11 +48,11 @@ class ArgSpec:
     help: str
     default: Any = None
     type: Any = None
-    action: Optional[str] = None
-    nargs: Optional[str] = None
+    action: str | None = None
+    nargs: str | None = None
 
 
-ARG_SPECS: Dict[str, ArgSpec] = {
+ARG_SPECS: dict[str, ArgSpec] = {
     "top_traj_file": ArgSpec(
         type=str,
         nargs="+",
@@ -145,6 +145,12 @@ class ArgSpec:
         default=True,
         help="Use bonded axes to rotate forces for UA level vibrational entropies",
     ),
+    "search_type": ArgSpec(
+        type=str,
+        default="RAD",
+        help="Type of neighbor search to use."
+        "Default RAD; grid search is also available",
+    ),
 }
 
 
@@ -159,7 +165,7 @@ class ConfigResolver:
       - validating trajectory-related numeric parameters
     """
 
-    def __init__(self, arg_specs: Optional[Dict[str, ArgSpec]] = None) -> None:
+    def __init__(self, arg_specs: dict[str, ArgSpec] | None = None) -> None:
         """Initialize the manager.
 
         Args:
@@ -168,7 +174,7 @@ def __init__(self, arg_specs: Optional[Dict[str, ArgSpec]] = None) -> None:
         """
         self._arg_specs = dict(arg_specs or ARG_SPECS)
 
-    def load_config(self, directory_path: str) -> Dict[str, Any]:
+    def load_config(self, directory_path: str) -> dict[str, Any]:
         """Load the first YAML config file found in a directory.
 
         The current behavior matches your existing workflow:
@@ -188,7 +194,7 @@ def load_config(self, directory_path: str) -> Dict[str, Any]:
 
         config_path = yaml_files[0]
         try:
-            with open(config_path, "r", encoding="utf-8") as file:
+            with open(config_path, encoding="utf-8") as file:
                 config = yaml.safe_load(file) or {"run1": {}}
             logger.info("Loaded configuration from: %s", config_path)
             return config
@@ -253,7 +259,7 @@ def build_parser(self) -> argparse.ArgumentParser:
                 )
                 continue
 
-            kwargs: Dict[str, Any] = {}
+            kwargs: dict[str, Any] = {}
             if spec.type is not None:
                 kwargs["type"] = spec.type
             if spec.default is not None:
@@ -266,7 +272,7 @@ def build_parser(self) -> argparse.ArgumentParser:
         return parser
 
     def resolve(
-        self, args: argparse.Namespace, run_config: Optional[Dict[str, Any]]
+        self, args: argparse.Namespace, run_config: dict[str, Any] | None
     ) -> argparse.Namespace:
         """Merge CLI arguments with YAML configuration and adjust logging level.
 
@@ -306,8 +312,8 @@ def resolve(
 
     @staticmethod
     def _detect_cli_overrides(
-        args_dict: Dict[str, Any], default_dict: Dict[str, Any]
-    ) -> Set[str]:
+        args_dict: dict[str, Any], default_dict: dict[str, Any]
+    ) -> set[str]:
         """Detect which args were explicitly overridden in the CLI.
 
         Args:
@@ -322,8 +328,8 @@ def _detect_cli_overrides(
     def _apply_yaml_defaults(
         self,
         args: argparse.Namespace,
-        run_config: Dict[str, Any],
-        cli_provided: Set[str],
+        run_config: dict[str, Any],
+        cli_provided: set[str],
     ) -> None:
         """Apply YAML values onto args for keys not provided by CLI.
 
@@ -336,7 +342,7 @@ def _apply_yaml_defaults(
             if yaml_value is None or key in cli_provided:
                 continue
             if key in self._arg_specs:
-                logger.debug("Using YAML value for %s: %s", key, yaml_value)
+                logger.debug(f"Using YAML value for {key}: {yaml_value}")
                 setattr(args, key, yaml_value)
 
     def _ensure_defaults(self, args: argparse.Namespace) -> None:

CodeEntropy/config/runtime.py

@@ -18,7 +18,7 @@
 import logging
 import os
 import pickle
-from typing import Any, Dict, Optional
+from typing import Any
 
 import MDAnalysis as mda
 import requests
@@ -112,7 +112,7 @@ def create_job_folder() -> str:
 
         return new_folder_path
 
-    def load_citation_data(self) -> Optional[Dict[str, Any]]:
+    def load_citation_data(self) -> dict[str, Any] | None:
         """Load CITATION.cff from GitHub.
 
         If the request fails (offline, blocked, etc.), returns None.
@@ -335,7 +335,7 @@ def _build_universe(
         kcal_units = args.kcal_force_units
 
         if forcefile is None:
-            logger.debug("Loading Universe with %s and %s", tprfile, trrfile)
+            logger.debug(f"Loading Universe with {tprfile} and {trrfile}")
             return mda.Universe(tprfile, trrfile, format=fileformat)
 
         return universe_operations.merge_forces(

CodeEntropy/core/logging.py

@@ -16,7 +16,6 @@
 
 import logging
 import os
-from typing import Dict, Optional
 
 from rich.console import Console
 from rich.logging import RichHandler
@@ -55,7 +54,7 @@ class LoggingConfig:
         handlers: Mapping of handler name to handler instance.
     """
 
-    _console: Optional[Console] = None
+    _console: Console | None = None
 
     @classmethod
     def get_console(cls) -> Console:
@@ -80,7 +79,7 @@ def __init__(self, folder: str, level: int = logging.INFO) -> None:
 
         self.level = level
         self.console = self.get_console()
-        self.handlers: Dict[str, logging.Handler] = {}
+        self.handlers: dict[str, logging.Handler] = {}
 
         self._setup_handlers()
 

CodeEntropy/entropy/configurational.py

@@ -10,46 +10,15 @@
 from __future__ import annotations
 
 import logging
-from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 
 import numpy as np
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass(frozen=True)
-class ConformationConfig:
-    """Configuration for assigning conformational states from a dihedral.
-
-    Attributes:
-        bin_width: Histogram bin width in degrees for peak detection.
-        start: Inclusive start frame index for trajectory slicing.
-        end: Exclusive end frame index for trajectory slicing.
-        step: Stride for trajectory slicing (must be positive).
-    """
-
-    bin_width: int
-    start: int
-    end: int
-    step: int
-
-
 class ConformationalEntropy:
-    """Assign dihedral conformational states and compute conformational entropy.
-
-    This class contains two independent responsibilities:
-      1) `assign_conformation`: Map a single dihedral angle time series to discrete
-         state labels by detecting histogram peaks and assigning the nearest peak.
-      2) `conformational_entropy_calculation`: Compute Shannon entropy of the
-         state distribution (in J/mol/K).
-
-    Notes:
-        `number_frames` is accepted by `conformational_entropy_calculation` for
-        compatibility with calling sites that track frame counts, but the entropy
-        is computed from the observed state counts (i.e., `len(states)`), which is
-        the correct normalization for the sampled distribution.
-    """
+    """Compute conformational entropy from states information."""
 
     _GAS_CONST: float = 8.3144598484848
 
@@ -61,69 +30,7 @@ def __init__(self) -> None:
         """
         pass
 
-    def assign_conformation(
-        self,
-        data_container: Any,
-        dihedral: Any,
-        number_frames: int,
-        bin_width: int,
-        start: int,
-        end: int,
-        step: int,
-    ) -> np.ndarray:
-        """Assign discrete conformational states for a single dihedral.
-
-        The dihedral angle time series is:
-          1) Collected across the trajectory slice [start:end:step].
-          2) Converted to [0, 360) degrees.
-          3) Histogrammed using `bin_width`.
-          4) Peaks are identified as bins with locally maximal population.
-          5) Each frame is assigned the index of the nearest peak.
-
-        Args:
-            data_container: MDAnalysis Universe/AtomGroup with a trajectory.
-            dihedral: Object providing `value()` for the current frame dihedral.
-            number_frames: Provided for call-site compatibility; not used for sizing.
-            bin_width: Histogram bin width in degrees.
-            start: Inclusive start frame index.
-            end: Exclusive end frame index.
-            step: Stride for trajectory slicing.
-
-        Returns:
-            Array of integer state labels of length equal to the trajectory slice.
-            Returns an empty array if the slice is empty.
-
-        Raises:
-            ValueError: If `bin_width` or `step` are invalid.
-        """
-        _ = number_frames
-
-        config = ConformationConfig(
-            bin_width=int(bin_width),
-            start=int(start),
-            end=int(end),
-            step=int(step),
-        )
-        self._validate_assignment_config(config)
-
-        traj_slice = data_container.trajectory[config.start : config.end : config.step]
-        n_slice = len(traj_slice)
-        if n_slice <= 0:
-            return np.array([], dtype=int)
-
-        phi = self._collect_dihedral_angles(traj_slice, dihedral)
-        peak_values = self._find_histogram_peaks(phi, config.bin_width)
-
-        if peak_values.size == 0:
-            return np.zeros(n_slice, dtype=int)
-
-        states = self._assign_nearest_peaks(phi, peak_values)
-        logger.debug("Final conformations: %s", states)
-        return states
-
-    def conformational_entropy_calculation(
-        self, states: Any, number_frames: int
-    ) -> float:
+    def conformational_entropy_calculation(self, states: Any) -> float:
         """Compute conformational entropy for a sequence of state labels.
 
         Entropy is computed as:
@@ -139,8 +46,6 @@ def conformational_entropy_calculation(
         Returns:
             float: Conformational entropy in J/mol/K.
         """
-        _ = number_frames
-
         arr = self._to_1d_array(states)
         if arr is None or arr.size == 0:
             return 0.0
@@ -154,96 +59,11 @@ def conformational_entropy_calculation(
         probs = probs[probs > 0.0]
 
         s_conf = -self._GAS_CONST * float(np.sum(probs * np.log(probs)))
-        logger.debug("Total conformational entropy: %s", s_conf)
+        logger.debug(f"Total conformational entropy: {s_conf}")
         return s_conf
 
     @staticmethod
-    def _validate_assignment_config(config: ConformationConfig) -> None:
-        """Validate conformation assignment configuration.
-
-        Args:
-            config: Assignment configuration.
-
-        Raises:
-            ValueError: If configuration values are invalid.
-        """
-        if config.step <= 0:
-            raise ValueError("step must be a positive integer")
-        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:
-            logger.warning(
-                "bin_width=%s does not evenly divide 360; histogram bins will be "
-                "uneven.",
-                config.bin_width,
-            )
-
-    @staticmethod
-    def _collect_dihedral_angles(traj_slice: Any, dihedral: Any) -> np.ndarray:
-        """Collect dihedral angles for each frame in the trajectory slice.
-
-        Args:
-            traj_slice: Slice of a trajectory iterable where iterating advances frames.
-            dihedral: Object with `value()` returning the dihedral in degrees.
-
-        Returns:
-            Array of dihedral values mapped into [0, 360).
-        """
-        phi = np.zeros(len(traj_slice), dtype=float)
-        for i, _ts in enumerate(traj_slice):
-            value = float(dihedral.value())
-            if value < 0.0:
-                value += 360.0
-            phi[i] = value
-        return phi
-
-    @staticmethod
-    def _find_histogram_peaks(phi: np.ndarray, bin_width: int) -> np.ndarray:
-        """Identify peak bin centers from a histogram of dihedral angles.
-
-        A peak is defined as a bin whose population is greater than or equal to
-        its immediate neighbors (with circular handling at the final bin).
-
-        Args:
-            phi: Dihedral angles in degrees, in [0, 360).
-            bin_width: Histogram bin width in degrees.
-
-        Returns:
-            1D array of peak bin center values (degrees). Empty if no peaks found.
-        """
-        number_bins = int(360 / bin_width)
-        popul, bin_edges = np.histogram(phi, bins=number_bins, range=(0.0, 360.0))
-        bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
-
-        peaks: list[float] = []
-        for idx in range(number_bins):
-            if popul[idx] == 0:
-                continue
-
-            left = popul[idx - 1] if idx > 0 else popul[number_bins - 1]
-            right = popul[idx + 1] if idx < number_bins - 1 else popul[0]
-
-            if popul[idx] >= left and popul[idx] >= right:
-                peaks.append(float(bin_centers[idx]))
-
-        return np.asarray(peaks, dtype=float)
-
-    @staticmethod
-    def _assign_nearest_peaks(phi: np.ndarray, peak_values: np.ndarray) -> np.ndarray:
-        """Assign each phi value to the index of its nearest peak.
-
-        Args:
-            phi: Dihedral angles in degrees.
-            peak_values: Peak centers (degrees).
-
-        Returns:
-            Integer state labels aligned with `phi`.
-        """
-        distances = np.abs(phi[:, None] - peak_values[None, :])
-        return np.argmin(distances, axis=1).astype(int)
-
-    @staticmethod
-    def _to_1d_array(states: Any) -> Optional[np.ndarray]:
+    def _to_1d_array(states: Any) -> np.ndarray | None:
         """Convert a state sequence into a 1D numpy array.
 
         Args: