Source code for CodeEntropy.levels.nodes.conformations

"""Compute conformational states for configurational entropy calculations.

This module defines a static DAG node that scans the selected trajectory frames
and builds conformational state descriptors (united-atom and residue level).
The resulting states are stored in ``shared_data`` for later use by
configurational entropy calculations.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from CodeEntropy.levels.dihedrals import ConformationStateBuilder
from CodeEntropy.trajectory.frames import FrameSelection

SharedData = dict[str, Any]
ConformationalStates = dict[str, Any]
FlexibleStates = dict[str, Any]




@dataclass(frozen=True)
class ConformationalStateConfig:
    """Configuration for conformational state construction.

    Attributes:
        n_frames: Number of frames to be analysed.
        bin_width: Histogram bin width in degrees.
    """

    n_frames: int
    bin_width: int





class ComputeConformationalStatesNode:
    """Static node that computes conformational states from trajectory dihedrals.

    Produces:
        shared_data["conformational_states"] = {"ua": states_ua, "res": states_res}
        shared_data["flexible_dihedrals"] = {"ua": flexible_ua, "res": flexible_res}

    Where:
        - states_ua is a dict keyed by ``(group_id, local_residue_id)``.
        - states_res is a list-like structure indexed by group id.
        - flexible_ua is a dict keyed by ``(group_id, local_residue_id)``.
        - flexible_res is a list-like structure indexed by group id.

    Notes:
        Frame selection is provided through ``shared_data["frame_selection"]``.
        During the current migration stage, that selection uses local
        analysis-universe frame indices because the workflow still physically
        frame-slices the universe.
    """

    def __init__(self, universe_operations: Any) -> None:
        """Initialize the node.

        Args:
            universe_operations: Object providing universe selection utilities used
                by ``ConformationStateBuilder``.
        """
        self._dihedral_analysis = ConformationStateBuilder(
            universe_operations=universe_operations
        )



    def run(
        self, shared_data: SharedData, *, progress: object | None = None
    ) -> dict[str, ConformationalStates]:
        """Compute conformational states and store them in shared_data.

        Args:
            shared_data: Shared data dictionary. Requires:
                - ``"reduced_universe"``
                - ``"levels"``
                - ``"groups"``
                - ``"frame_selection"``
                - ``"args"`` with attribute ``bin_width``
            progress: Optional progress sink provided by ResultsReporter.progress().

        Returns:
            Dict containing ``"conformational_states"``.
        """
        u = shared_data["reduced_universe"]
        levels = shared_data["levels"]
        groups = shared_data["groups"]
        frame_selection: FrameSelection = shared_data["frame_selection"]
        bin_width = int(shared_data["args"].bin_width)

        states_ua, states_res, flexible_ua, flexible_res = (
            self._dihedral_analysis.build_conformational_states(
                data_container=u,
                levels=levels,
                groups=groups,
                bin_width=bin_width,
                frame_selection=frame_selection,
                progress=progress,
            )
        )

        conformational_states: ConformationalStates = {
            "ua": states_ua,
            "res": states_res,
        }
        shared_data["conformational_states"] = conformational_states

        flexible_states: FlexibleStates = {
            "ua": flexible_ua,
            "res": flexible_res,
        }
        shared_data["flexible_dihedrals"] = flexible_states

        return {"conformational_states": conformational_states}