Merge pull request #366 from CCPBioSim/362-introduce-a-dedicated-conformationdag-stage

Introduce Dedicated ConformationDAG Stage for Conformational Analysis

Author: harryswift01

CodeEntropy/levels/nodes/conformations.py CodeEntropy/levels/conformation_dag.pyCodeEntropy/levels/nodes/conformations.py renamed to CodeEntropy/levels/conformation_dag.py

@@ -1,4 +1,8 @@
-"""Compute conformational states for configurational entropy calculations."""
+"""Conformational-state DAG orchestration.
+
+This module owns the conformational stage between static structural setup and
+frame-local covariance/neighbour execution.
+"""
 
 from __future__ import annotations
 
@@ -12,26 +16,23 @@
 FlexibleStates = dict[str, Any]
 
 
-class ComputeConformationalStatesNode:
-    """Static node that computes conformational states from selected frames.
-
-    Produces:
-        shared_data["conformational_states"] = {"ua": states_ua, "res": states_res}
-        shared_data["flexible_dihedrals"] = {"ua": flexible_ua, "res": flexible_res}
-    """
+class ConformationDAG:
+    """Execute conformational-state construction for selected trajectory frames."""
 
     def __init__(self, universe_operations: Any | None = None) -> None:
-        """Initialise the conformational-state node.
-
-        Args:
-            universe_operations: Optional universe-operation adapter passed to the
-                underlying conformation-state builder.
-        """
         self._builder = ConformationStateBuilder(
             universe_operations=universe_operations
         )
 
-    def run(
+    def build(self) -> ConformationDAG:
+        """Build the conformational DAG topology.
+
+        Returns:
+            Self, to allow fluent construction.
+        """
+        return self
+
+    def execute(
         self,
         shared_data: SharedData,
         *,
@@ -43,12 +44,8 @@ def run(
             shared_data: Shared workflow data containing ``reduced_universe``,
                 ``levels``, ``groups``, ``frame_selection``, and ``args.bin_width``.
             progress: Optional progress sink forwarded to the conformation builder.
-
         Returns:
             A dictionary containing the computed ``conformational_states`` mapping.
-
-        Raises:
-            KeyError: If required entries are missing from ``shared_data``.
         """
         universe = shared_data["reduced_universe"]
         levels = shared_data["levels"]

CodeEntropy/levels/level_dag.py

@@ -1,8 +1,8 @@
 """Hierarchy-level DAG orchestration.
 
 LevelDAG owns hierarchy-level workflow order. Static setup nodes prepare
-structural and conformational data, then frame-local covariance and neighbour
-observables are executed through deterministic frame map-reduce.
+structural data. ConformationDAG computes trajectory-series conformational
+states. FrameScheduler executes frame-local covariance and neighbour work.
 """
 
 from __future__ import annotations
@@ -12,14 +12,14 @@
 import networkx as nx
 
 from CodeEntropy.levels.axes import AxesCalculator
+from CodeEntropy.levels.conformation_dag import ConformationDAG
 from CodeEntropy.levels.execution.policy import ExecutionPolicy
 from CodeEntropy.levels.execution.reducers import NeighborReducer
 from CodeEntropy.levels.execution.scheduler import FrameScheduler
 from CodeEntropy.levels.frame_dag import FrameGraph
 from CodeEntropy.levels.neighbors import Neighbors
 from CodeEntropy.levels.nodes.accumulators import InitCovarianceAccumulatorsNode
 from CodeEntropy.levels.nodes.beads import BuildBeadsNode
-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.results.reporter import _RichProgressSink
@@ -38,15 +38,14 @@ def __init__(self, universe_operations: Any | None = None) -> None:
         self._universe_operations = universe_operations
         self._static_graph = nx.DiGraph()
         self._static_nodes: dict[str, Any] = {}
+        self._conformation_dag = ConformationDAG(
+            universe_operations=universe_operations
+        )
         self._frame_dag = FrameGraph(universe_operations=universe_operations)
         self._policy = ExecutionPolicy()
 
     def build(self) -> LevelDAG:
-        """Build static and frame-level DAG topology.
-
-        Returns:
-            The current ``LevelDAG`` instance for fluent construction.
-        """
+        """Build the static, conformation, and frame DAG topology."""
         self._add_static("detect_molecules", DetectMoleculesNode())
         self._add_static("detect_levels", DetectLevelsNode(), deps=["detect_molecules"])
         self._add_static("build_beads", BuildBeadsNode(), deps=["detect_levels"])
@@ -55,12 +54,8 @@ def build(self) -> LevelDAG:
             InitCovarianceAccumulatorsNode(),
             deps=["detect_levels"],
         )
-        self._add_static(
-            "compute_conformational_states",
-            ComputeConformationalStatesNode(self._universe_operations),
-            deps=["detect_levels"],
-        )
 
+        self._conformation_dag.build()
         self._frame_dag.build()
         return self
 
@@ -87,6 +82,8 @@ def execute(
         shared_data.setdefault("axes_manager", AxesCalculator())
 
         self._run_static_stage(shared_data, progress=progress)
+        self._run_conformation_stage(shared_data, progress=progress)
+
         self._initialise_neighbor_metadata(shared_data)
         NeighborReducer.initialise(shared_data)
         self._run_frame_stage(shared_data, progress=progress)
@@ -137,6 +134,15 @@ def _add_static(
         for dep in deps or []:
             self._static_graph.add_edge(dep, name)
 
+    def _run_conformation_stage(
+        self,
+        shared_data: dict[str, Any],
+        *,
+        progress: _RichProgressSink | None = None,
+    ) -> None:
+        """Run conformational-state construction after static setup."""
+        self._conformation_dag.execute(shared_data, progress=progress)
+
     def _run_frame_stage(
         self,
         shared_data: dict[str, Any],

docs/api/CodeEntropy.levels.conformation_dag.rst

@@ -0,0 +1,7 @@
+CodeEntropy.levels.conformation\_dag module
+===========================================
+
+.. automodule:: CodeEntropy.levels.conformation_dag
+   :members:
+   :show-inheritance:
+   :undoc-members:

docs/api/CodeEntropy.levels.nodes.conformations.rst

@@ -14,7 +14,6 @@ Submodules
 
    CodeEntropy.levels.nodes.accumulators
    CodeEntropy.levels.nodes.beads
-   CodeEntropy.levels.nodes.conformations
    CodeEntropy.levels.nodes.covariance
    CodeEntropy.levels.nodes.detect_levels
    CodeEntropy.levels.nodes.detect_molecules

docs/api/CodeEntropy.levels.nodes.rst

@@ -21,6 +21,7 @@ Submodules
    :maxdepth: 4
 
    CodeEntropy.levels.axes
+   CodeEntropy.levels.conformation_dag
    CodeEntropy.levels.dihedrals
    CodeEntropy.levels.forces
    CodeEntropy.levels.frame_dag

docs/api/CodeEntropy.levels.rst

@@ -1,11 +1,11 @@
-"""Unit tests for the conformational-state static node."""
+"""Unit tests for the conformational-state DAG stage."""
 
 from __future__ import annotations
 
 from types import SimpleNamespace
 
-from CodeEntropy.levels.nodes import conformations
-from CodeEntropy.levels.nodes.conformations import ComputeConformationalStatesNode
+from CodeEntropy.levels import conformation_dag
+from CodeEntropy.levels.conformation_dag import ConformationDAG
 
 
 class FakeConformationStateBuilder:
@@ -43,7 +43,13 @@ def build_conformational_states(
         )
 
 
-def test_compute_conformational_states_node_runs_and_writes_shared_data(monkeypatch):
+def test_conformation_dag_build_returns_self():
+    dag = ConformationDAG()
+
+    assert dag.build() is dag
+
+
+def test_conformation_dag_executes_builder_and_writes_shared_data(monkeypatch):
     builder_holder = {}
 
     def builder_factory(universe_operations):
@@ -52,13 +58,13 @@ def builder_factory(universe_operations):
         return builder
 
     monkeypatch.setattr(
-        conformations,
+        conformation_dag,
         "ConformationStateBuilder",
         builder_factory,
     )
 
     universe_operations = object()
-    node = ComputeConformationalStatesNode(universe_operations)
+    dag = ConformationDAG(universe_operations=universe_operations)
 
     universe = object()
     frame_selection = object()
@@ -72,7 +78,7 @@ def builder_factory(universe_operations):
         "args": SimpleNamespace(bin_width=30),
     }
 
-    result = node.run(shared_data, progress=progress)
+    result = dag.execute(shared_data, progress=progress)
 
     assert shared_data["conformational_states"] == {
         "ua": {"ua_key": ["state_a"]},
@@ -100,20 +106,24 @@ def builder_factory(universe_operations):
     ]
 
 
-def test_compute_conformational_states_node_converts_bin_width_to_int(monkeypatch):
+def test_conformation_dag_converts_bin_width_to_int(monkeypatch):
     captured = {}
 
     class Builder:
         def __init__(self, universe_operations):
-            pass
+            self.universe_operations = universe_operations
 
         def build_conformational_states(self, **kwargs):
             captured.update(kwargs)
             return {}, [], {}, []
 
-    monkeypatch.setattr(conformations, "ConformationStateBuilder", Builder)
+    monkeypatch.setattr(
+        conformation_dag,
+        "ConformationStateBuilder",
+        Builder,
+    )
 
-    node = ComputeConformationalStatesNode()
+    dag = ConformationDAG()
     shared_data = {
         "reduced_universe": object(),
         "levels": [],
@@ -122,6 +132,6 @@ def build_conformational_states(self, **kwargs):
         "args": SimpleNamespace(bin_width="45"),
     }
 
-    node.run(shared_data)
+    dag.execute(shared_data)
 
     assert captured["bin_width"] == 45