CodeEntropy 2.0.0__py3-none-any.whl → 2.1.1__py3-none-any.whl

CodeEntropy/__init__.py

@@ -8,4 +8,4 @@ entropy estimates, supporting a wide range of applications in molecular simulati
 and statistical mechanics.
 """
 
-__version__ = "2.0.0"
+__version__ = "2.1.1"

CodeEntropy/config/argparse.py

@@ -25,7 +25,7 @@ import glob
 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 @@ ARG_SPECS: Dict[str, 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 @@ class ConfigResolver:
         """
         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 @@ class ConfigResolver:
 
         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 @@ class ConfigResolver:
                 )
                 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 @@ class ConfigResolver:
         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 @@ class ConfigResolver:
 
     @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 @@ class ConfigResolver:
     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 @@ class ConfigResolver:
             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 @@ from __future__ import annotations
 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 @@ class CodeEntropyRunner:
 
         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 @@ class CodeEntropyRunner:
         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 @@ from __future__ import annotations
 
 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 @@ class LoggingConfig:
 
         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 @@ This module provides:
 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 @@ class ConformationalEntropy:
         """
         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 @@ class ConformationalEntropy:
         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 @@ class ConformationalEntropy:
         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:

CodeEntropy/entropy/graph.py

@@ -15,18 +15,19 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass
-from typing import Any, Dict
+from typing import Any
 
 import networkx as nx
 
 from CodeEntropy.entropy.nodes.aggregate import AggregateEntropyNode
 from CodeEntropy.entropy.nodes.configurational import ConfigurationalEntropyNode
+from CodeEntropy.entropy.nodes.orientational import OrientationalEntropyNode
 from CodeEntropy.entropy.nodes.vibrational import VibrationalEntropyNode
 
 logger = logging.getLogger(__name__)
 
 
-SharedData = Dict[str, Any]
+SharedData = dict[str, Any]
 
 
 @dataclass(frozen=True)
@@ -57,9 +58,9 @@ class EntropyGraph:
     def __init__(self) -> None:
         """Initialize an empty entropy graph."""
         self._graph: nx.DiGraph = nx.DiGraph()
-        self._nodes: Dict[str, Any] = {}
+        self._nodes: dict[str, Any] = {}
 
-    def build(self) -> "EntropyGraph":
+    def build(self) -> EntropyGraph:
         """Populate the graph with the standard entropy workflow.
 
         Returns:
@@ -68,10 +69,15 @@ class EntropyGraph:
         specs = (
             NodeSpec("vibrational_entropy", VibrationalEntropyNode()),
             NodeSpec("configurational_entropy", ConfigurationalEntropyNode()),
+            NodeSpec("orientational_entropy", OrientationalEntropyNode()),
             NodeSpec(
                 "aggregate_entropy",
                 AggregateEntropyNode(),
-                deps=("vibrational_entropy", "configurational_entropy"),
+                deps=(
+                    "vibrational_entropy",
+                    "configurational_entropy",
+                    "orientational_entropy",
+                ),
             ),
         )
 
@@ -82,7 +88,7 @@ class EntropyGraph:
 
     def execute(
         self, shared_data: SharedData, *, progress: object | None = None
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """Execute the entropy graph in topological order.
 
         Nodes are executed in dependency order (topological sort). Each node reads
@@ -105,7 +111,7 @@ class EntropyGraph:
         Raises:
             KeyError: If a node name is missing from the internal node registry.
         """
-        results: Dict[str, Any] = {}
+        results: dict[str, Any] = {}
 
         for node_name in nx.topological_sort(self._graph):
             node = self._nodes[node_name]

CodeEntropy/entropy/nodes/aggregate.py

@@ -2,10 +2,11 @@
 
 from __future__ import annotations
 
+from collections.abc import Mapping, MutableMapping
 from dataclasses import dataclass
-from typing import Any, Dict, Mapping, MutableMapping, Optional
+from typing import Any
 
-EntropyResults = Dict[str, Any]
+EntropyResults = dict[str, Any]
 
 
 @dataclass(frozen=True, slots=True)
@@ -25,11 +26,12 @@ class AggregateEntropyNode:
 
     vibrational_key: str = "vibrational_entropy"
     configurational_key: str = "configurational_entropy"
+    orientational_key: str = "orientational_entropy"
     output_key: str = "entropy_results"
 
     def run(
         self, shared_data: MutableMapping[str, Any], **_: Any
-    ) -> Dict[str, EntropyResults]:
+    ) -> dict[str, EntropyResults]:
         """Run the aggregation step.
 
         Args:
@@ -69,10 +71,13 @@ class AggregateEntropyNode:
             "configurational_entropy": self._get_optional(
                 shared_data, self.configurational_key
             ),
+            "orientational_entropy": self._get_optional(
+                shared_data, self.orientational_key
+            ),
         }
 
     @staticmethod
-    def _get_optional(shared_data: Mapping[str, Any], key: str) -> Optional[Any]:
+    def _get_optional(shared_data: Mapping[str, Any], key: str) -> Any | None:
         """Fetch an optional value from shared data.
 
         Args:

CodeEntropy/entropy/nodes/configurational.py

@@ -3,16 +3,9 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Iterable, Mapping, MutableMapping, Sequence
 from typing import (
     Any,
-    Dict,
-    Iterable,
-    Mapping,
-    MutableMapping,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
 )
 
 import numpy as np
@@ -23,8 +16,8 @@ logger = logging.getLogger(__name__)
 
 GroupId = int
 ResidueId = int
-StateKey = Tuple[GroupId, ResidueId]
-StateSequence = Union[Sequence[Any], np.ndarray]
+StateKey = tuple[GroupId, ResidueId]
+StateSequence = Sequence[Any] | np.ndarray
 
 
 class ConfigurationalEntropyNode:
@@ -36,7 +29,7 @@ class ConfigurationalEntropyNode:
     Results are written back into ``shared_data["configurational_entropy"]``.
     """
 
-    def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> Dict[str, Any]:
+    def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> dict[str, Any]:
         """Execute configurational entropy calculation.
 
         Args:
@@ -58,7 +51,7 @@ class ConfigurationalEntropyNode:
         ce = self._build_entropy_engine()
 
         fragments = universe.atoms.fragments
-        results: Dict[int, Dict[str, float]] = {}
+        results: dict[int, dict[str, float]] = {}
 
         for group_id, mol_ids in groups.items():
             results[group_id] = {"ua": 0.0, "res": 0.0, "poly": 0.0}
@@ -104,9 +97,9 @@ class ConfigurationalEntropyNode:
 
     def _get_state_containers(
         self, shared_data: Mapping[str, Any]
-    ) -> Tuple[
-        Dict[StateKey, StateSequence],
-        Union[Dict[GroupId, StateSequence], Sequence[Optional[StateSequence]]],
+    ) -> tuple[
+        dict[StateKey, StateSequence],
+        dict[GroupId, StateSequence] | Sequence[StateSequence | None],
     ]:
         """Retrieve conformational state containers.
 
@@ -144,7 +137,7 @@ class ConfigurationalEntropyNode:
         residues: Iterable[Any],
         states_ua: Mapping[StateKey, StateSequence],
         n_frames: int,
-        reporter: Optional[Any],
+        reporter: Any | None,
     ) -> float:
         """Compute united atom entropy for a group.
 
@@ -163,7 +156,7 @@ class ConfigurationalEntropyNode:
 
         for res_id, res in enumerate(residues):
             states = states_ua.get((group_id, res_id))
-            val = self._entropy_or_zero(ce, states, n_frames)
+            val = self._entropy_or_zero(ce, states)
             total += val
 
             if reporter is not None:
@@ -186,29 +179,28 @@ class ConfigurationalEntropyNode:
         *,
         ce: ConformationalEntropy,
         group_id: int,
-        states_res: Union[Dict[int, StateSequence], Sequence[Optional[StateSequence]]],
+        states_res: dict[int, StateSequence] | Sequence[StateSequence | None],
         n_frames: int,
     ) -> float:
         """Compute residue-level entropy for a group."""
         group_states = self._get_group_states(states_res, group_id)
-        return self._entropy_or_zero(ce, group_states, n_frames)
+        return self._entropy_or_zero(ce, group_states)
 
     def _entropy_or_zero(
         self,
         ce: ConformationalEntropy,
-        states: Optional[StateSequence],
-        n_frames: int,
+        states: StateSequence | None,
     ) -> float:
         """Return entropy value or zero if no state data exists."""
         if not self._has_state_data(states):
             return 0.0
-        return float(ce.conformational_entropy_calculation(states, n_frames))
+        return float(ce.conformational_entropy_calculation(states))
 
     @staticmethod
     def _get_group_states(
-        states_res: Union[Dict[int, StateSequence], Sequence[Optional[StateSequence]]],
+        states_res: dict[int, StateSequence] | Sequence[StateSequence | None],
         group_id: int,
-    ) -> Optional[StateSequence]:
+    ) -> StateSequence | None:
         """Fetch group states from container."""
         if isinstance(states_res, dict):
             return states_res.get(group_id)
@@ -217,7 +209,7 @@ class ConfigurationalEntropyNode:
         return None
 
     @staticmethod
-    def _has_state_data(states: Optional[StateSequence]) -> bool:
+    def _has_state_data(states: StateSequence | None) -> bool:
         """Check if state container has usable data."""
         if states is None:
             return False