Introduce Basic Orientational Entropy Calculations

CodeEntropy/config/argparse.py

@@ -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",
+    ),
 }
 
 

CodeEntropy/entropy/configurational.py

@@ -223,7 +223,7 @@ def _find_histogram_peaks(phi: np.ndarray, bin_width: int) -> np.ndarray:
             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:
+            if popul[idx] >= left and popul[idx] > right:
                 peaks.append(float(bin_centers[idx]))
 
         return np.asarray(peaks, dtype=float)

CodeEntropy/entropy/graph.py

@@ -21,6 +21,7 @@
 
 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__)
@@ -68,10 +69,15 @@ def build(self) -> "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",
+                ),
             ),
         )
 

CodeEntropy/entropy/nodes/aggregate.py

@@ -25,6 +25,7 @@ class AggregateEntropyNode:
 
     vibrational_key: str = "vibrational_entropy"
     configurational_key: str = "configurational_entropy"
+    orientational_key: str = "orientational_entropy"
     output_key: str = "entropy_results"
 
     def run(
@@ -69,6 +70,9 @@ def _collect_entropy_results(
             "configurational_entropy": self._get_optional(
                 shared_data, self.configurational_key
             ),
+            "orientational_entropy": self._get_optional(
+                shared_data, self.orientational_key
+            ),
         }
 
     @staticmethod

CodeEntropy/entropy/nodes/orientational.py

@@ -0,0 +1,85 @@
+"""Node for computing orientational entropy from neighbors."""
+
+from __future__ import annotations
+
+import logging
+from typing import (
+    Any,
+    Dict,
+    MutableMapping,
+    Sequence,
+    Tuple,
+    Union,
+)
+
+import numpy as np
+
+from CodeEntropy.entropy.orientational import OrientationalEntropy
+
+logger = logging.getLogger(__name__)
+
+GroupId = int
+ResidueId = int
+StateKey = Tuple[GroupId, ResidueId]
+StateSequence = Union[Sequence[Any], np.ndarray]
+
+
+class OrientationalEntropyNode:
+    """Compute orientational entropy using precomputed neighbors and symmetry.
+
+    This node reads number of neighbors and symmetry from ``shared_data`` and
+    computes entropy contributions at the molecular (highest) level.
+
+    Results are written back into ``shared_data["orientational_entropy"]``.
+    """
+
+    def run(self, shared_data: MutableMapping[str, Any], **_: Any) -> Dict[str, Any]:
+        """Execute orientational entropy calculation.
+
+        Args:
+            shared_data: Shared workflow state dictionary.
+
+        Returns:
+            Dictionary containing orientational entropy results.
+
+        Raises:
+            KeyError: If required keys are missing.
+        """
+        groups = shared_data["groups"]
+        levels = shared_data["levels"]
+        neighbors = shared_data["neighbors"]
+        symmetry_number = shared_data["symmetry_number"]
+        linear = shared_data["linear"]
+        reporter = shared_data.get("reporter")
+
+        oe = self._build_entropy_engine()
+
+        results: Dict[int, float] = {}
+
+        for group_id, mol_ids in groups.items():
+            rep_mol_id = mol_ids[0]
+            highest_level = levels[rep_mol_id][-1]
+
+            neighbor = neighbors[group_id]
+            symmetry = symmetry_number[group_id]
+            line = linear[group_id]
+
+            result_value = oe.calculate_orientational(
+                neighbor,
+                symmetry,
+                line,
+            )
+            results[group_id] = result_value
+
+            if reporter is not None:
+                reporter.add_results_data(
+                    group_id, highest_level, "Orientational", result_value
+                )
+
+        shared_data["orientational_entropy"] = results
+
+        return {"orientational_entropy": results}
+
+    def _build_entropy_engine(self) -> OrientationalEntropy:
+        """Create the entropy calculation engine."""
+        return OrientationalEntropy()

CodeEntropy/entropy/orientational.py

@@ -1,19 +1,14 @@
 """Orientational entropy calculations.
 
 This module defines `OrientationalEntropy`, which computes orientational entropy
-from a neighbor-count mapping.
-
-The current implementation supports non-water neighbors. Water-specific behavior
-can be implemented later behind an interface so the core calculation remains
-stable and testable.
+from a neighbour count and symmetry information.
 """
 
 from __future__ import annotations
 
 import logging
 import math
 from dataclasses import dataclass
-from typing import Any, Mapping
 
 import numpy as np
 
@@ -48,111 +43,77 @@ class OrientationalEntropy:
 
     def __init__(
         self,
-        run_manager: Any,
-        args: Any,
-        universe: Any,
-        reporter: Any,
-        group_molecules: Any,
         gas_constant: float = _GAS_CONST_J_PER_MOL_K,
     ) -> None:
         """Initialize the orientational entropy calculator.
 
         Args:
-            run_manager: Run manager (currently unused by this class).
-            args: User arguments (currently unused by this class).
-            universe: MDAnalysis Universe (currently unused by this class).
-            reporter: Data logger (currently unused by this class).
-            group_molecules: Grouping helper (currently unused by this class).
             gas_constant: Gas constant in J/(mol*K).
         """
-        self._run_manager = run_manager
-        self._args = args
-        self._universe = universe
-        self._reporter = reporter
-        self._group_molecules = group_molecules
         self._gas_constant = float(gas_constant)
 
-    def calculate(self, neighbours: Mapping[str, int]) -> OrientationalEntropyResult:
-        """Calculate orientational entropy from neighbor counts.
-
-        For each neighbor species (except water), the number of orientations is
-        estimated as:
+    def calculate_orientational(
+        self,
+        neighbour_count: float,
+        symmetry_number: int,
+        linear: bool,
+    ) -> OrientationalEntropyResult:
+        """Calculate orientational entropy from neighbour counts.
 
-            Ω = sqrt(Nc^3 * π)
+        The number of orientations is estimated as:
+            Ω = sqrt(N_av^3 * π)/symmetry_number for non-linear molecules
+            Ω = N_av / symmetry_number for linear molecules
 
         and the entropy contribution is:
 
             S = R * ln(Ω)
 
-        where Nc is the neighbor count and R is the gas constant.
+        where N_av is the average number of neighbours and R is the gas constant.
 
         Args:
-            neighbours: Mapping of neighbor species name to count.
+            neighbours: average number of neighbours
+            symmetry_number: symmetry number of molecule of interest
+            linear: True if molecule of interest is linear
 
         Returns:
             OrientationalEntropyResult containing the total entropy in J/mol/K.
-        """
-        total = 0.0
-        for species, count in neighbours.items():
-            if self._is_water(species):
-                logger.debug(
-                    "Skipping water species %s in orientational entropy.", species
-                )
-                continue
-
-            contribution = self._entropy_contribution(count)
-            logger.debug(
-                "Orientational entropy contribution for %s: %s", species, contribution
-            )
-            total += contribution
-
-        logger.debug("Final orientational entropy total: %s", total)
-        return OrientationalEntropyResult(total=float(total))
-
-    @staticmethod
-    def _is_water(species: str) -> bool:
-        """Return True if the species should be treated as water.
-
-        Args:
-            species: Species identifier.
-
-        Returns:
-            True if the species is considered water.
-        """
-        return species in {"H2O", "WAT", "HOH"}
-
-    def _entropy_contribution(self, neighbour_count: int) -> float:
-        """Compute the entropy contribution for a single neighbor count.
-
-        Args:
-            neighbour_count: Number of neighbors (Nc).
-
-        Returns:
-            Entropy contribution in J/mol/K.
 
         Raises:
-            ValueError: If neighbour_count is negative.
+            ValueError if number of neighbours is negative.
         """
         if neighbour_count < 0:
             raise ValueError(f"neighbour_count must be >= 0, got {neighbour_count}")
 
-        if neighbour_count == 0:
-            return 0.0
+        omega = self._omega(neighbour_count, symmetry_number, linear)
 
-        omega = self._omega(neighbour_count)
-        if omega <= 0.0:
-            return 0.0
+        total = self._gas_constant * math.log(omega)
+        logger.debug("Orientational entropy total: %s", total)
 
-        return self._gas_constant * math.log(omega)
+        return total
 
     @staticmethod
-    def _omega(neighbour_count: int) -> float:
+    def _omega(neighbour_count: int, symmetry: int, linear: bool) -> float:
         """Compute the number of orientations Ω.
 
         Args:
-            neighbour_count: Number of neighbors (Nc).
+            neighbour_count: average number of neighbours.
+            symmetry_number: The symmetry number of the molecule.
+            linear: Is the molecule linear (True or False).
 
         Returns:
             Ω (unitless).
         """
-        return float(np.sqrt((neighbour_count**3) * math.pi))
+        # symmetry number 0 = spherically symmetric = no orientational entropy
+        if symmetry == 0:
+            omega = 1
+        else:
+            if linear:
+                omega = neighbour_count / symmetry
+            else:
+                omega = np.sqrt((neighbour_count**3) * math.pi) / symmetry
+
+        # avoid negative orientational entropy
+        if omega < 1:
+            omega = 1
+
+        return omega

CodeEntropy/levels/dihedrals.py

@@ -362,13 +362,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])
 

CodeEntropy/levels/level_dag.py

@@ -28,6 +28,7 @@
 from CodeEntropy.levels.nodes.conformations import ComputeConformationalStatesNode
 from CodeEntropy.levels.nodes.detect_levels import DetectLevelsNode
 from CodeEntropy.levels.nodes.detect_molecules import DetectMoleculesNode
+from CodeEntropy.levels.nodes.find_neighbors import ComputeNeighborsNode
 
 logger = logging.getLogger(__name__)
 
@@ -81,6 +82,9 @@ def build(self) -> "LevelDAG":
             ComputeConformationalStatesNode(self._universe_operations),
             deps=["detect_levels"],
         )
+        self._add_static(
+            "find_neighbors", ComputeNeighborsNode(), deps=["detect_levels"]
+        )
 
         self._frame_dag.build()
         return self