refactor(levels): introduce dedicated ConformationDAG stage

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,44 +16,34 @@
 FlexibleStates = dict[str, Any]
 
 
-class ComputeConformationalStatesNode:
-    """Static node that computes conformational states from selected frames.
+class ConformationDAG:
+    """Execute conformational-state construction for selected trajectory frames.
 
-    Produces:
-        shared_data["conformational_states"] = {"ua": states_ua, "res": states_res}
-        shared_data["flexible_dihedrals"] = {"ua": flexible_ua, "res": flexible_res}
+    The first implementation intentionally preserves the existing serial
+    ConformationStateBuilder behaviour. Later issues can replace this internal
+    implementation with chunked map-reduce execution.
     """
 
     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,
         *,
         progress: object | None = None,
     ) -> dict[str, ConformationalStates]:
-        """Compute conformational states and store them in shared workflow data.
-
-        Args:
-            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``.
-        """
+        """Compute conformational states and store them in shared workflow data."""
         universe = shared_data["reduced_universe"]
         levels = shared_data["levels"]
         groups = shared_data["groups"]

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],

…/levels/nodes/test_conformations_node.py …eEntropy/levels/test_conformation_dag.pytests/unit/CodeEntropy/levels/nodes/test_conformations_node.py renamed to tests/unit/CodeEntropy/levels/test_conformation_dag.py tests/unit/CodeEntropy/levels/nodes/test_conformations_node.py renamed to tests/unit/CodeEntropy/levels/test_conformation_dag.py

@@ -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

tests/unit/CodeEntropy/levels/test_level_dag.py

@@ -7,16 +7,17 @@
 from CodeEntropy.levels.level_dag import LevelDAG
 
 
-def test_build_registers_static_nodes_and_builds_frame_dag():
+def test_build_registers_static_nodes_and_builds_stage_dags():
     with (
         patch("CodeEntropy.levels.level_dag.DetectMoleculesNode"),
         patch("CodeEntropy.levels.level_dag.DetectLevelsNode"),
         patch("CodeEntropy.levels.level_dag.BuildBeadsNode"),
         patch("CodeEntropy.levels.level_dag.InitCovarianceAccumulatorsNode"),
-        patch("CodeEntropy.levels.level_dag.ComputeConformationalStatesNode"),
+        patch("CodeEntropy.levels.level_dag.ConformationDAG"),
     ):
         universe_operations = MagicMock()
         dag = LevelDAG(universe_operations=universe_operations)
+        dag._conformation_dag.build = MagicMock()
         dag._frame_dag.build = MagicMock()
 
         out = dag.build()
@@ -27,15 +28,19 @@ def test_build_registers_static_nodes_and_builds_frame_dag():
         "detect_levels",
         "build_beads",
         "init_covariance_accumulators",
-        "compute_conformational_states",
     }
     assert "find_neighbors" not in dag._static_nodes
+    assert "compute_conformational_states" not in dag._static_nodes
 
     assert ("detect_molecules", "detect_levels") in dag._static_graph.edges
     assert ("detect_levels", "build_beads") in dag._static_graph.edges
     assert ("detect_levels", "init_covariance_accumulators") in dag._static_graph.edges
-    assert ("detect_levels", "compute_conformational_states") in dag._static_graph.edges
+    assert (
+        "detect_levels",
+        "compute_conformational_states",
+    ) not in dag._static_graph.edges
 
+    dag._conformation_dag.build.assert_called_once()
     dag._frame_dag.build.assert_called_once()
 
 
@@ -46,6 +51,7 @@ def test_execute_sets_default_axes_manager_and_runs_workflow_stages():
     progress = MagicMock()
 
     dag._run_static_stage = MagicMock()
+    dag._run_conformation_stage = MagicMock()
     dag._initialise_neighbor_metadata = MagicMock()
     dag._run_frame_stage = MagicMock()
 
@@ -59,6 +65,10 @@ def test_execute_sets_default_axes_manager_and_runs_workflow_stages():
     assert "axes_manager" in shared_data
 
     dag._run_static_stage.assert_called_once_with(shared_data, progress=progress)
+    dag._run_conformation_stage.assert_called_once_with(
+        shared_data,
+        progress=progress,
+    )
     dag._initialise_neighbor_metadata.assert_called_once_with(shared_data)
     initialise.assert_called_once_with(shared_data)
     dag._run_frame_stage.assert_called_once_with(shared_data, progress=progress)
@@ -126,6 +136,21 @@ def test_run_static_stage_falls_back_when_node_does_not_accept_progress():
     ]
 
 
+def test_run_conformation_stage_delegates_to_conformation_dag():
+    dag = LevelDAG()
+    shared_data = {}
+    progress = MagicMock()
+
+    dag._conformation_dag.execute = MagicMock()
+
+    dag._run_conformation_stage(shared_data, progress=progress)
+
+    dag._conformation_dag.execute.assert_called_once_with(
+        shared_data,
+        progress=progress,
+    )
+
+
 def test_run_frame_stage_collects_frame_indices_and_delegates_to_scheduler():
     universe_operations = MagicMock()
     dag = LevelDAG(universe_operations=universe_operations)
@@ -183,3 +208,49 @@ def test_initialise_neighbor_metadata_falls_back_to_universe_key():
         LevelDAG._initialise_neighbor_metadata(shared_data)
 
     helper.get_symmetry.assert_called_once_with(universe=universe, groups={0: [0]})
+
+
+def test_level_dag_runs_static_conformation_then_frame(monkeypatch):
+    dag = LevelDAG(universe_operations=object())
+    calls = []
+
+    monkeypatch.setattr(
+        dag,
+        "_run_static_stage",
+        lambda shared_data, progress=None: calls.append("static"),
+    )
+    monkeypatch.setattr(
+        dag,
+        "_run_conformation_stage",
+        lambda shared_data, progress=None: calls.append("conformation"),
+    )
+    monkeypatch.setattr(
+        dag,
+        "_initialise_neighbor_metadata",
+        lambda shared_data: calls.append("neighbor_metadata"),
+    )
+    monkeypatch.setattr(
+        dag,
+        "_run_frame_stage",
+        lambda shared_data, progress=None: calls.append("frame"),
+    )
+
+    monkeypatch.setattr(
+        "CodeEntropy.levels.level_dag.NeighborReducer.initialise",
+        lambda shared_data: calls.append("neighbor_initialise"),
+    )
+    monkeypatch.setattr(
+        "CodeEntropy.levels.level_dag.NeighborReducer.finalise",
+        lambda shared_data: calls.append("neighbor_finalise"),
+    )
+
+    dag.execute({})
+
+    assert calls == [
+        "static",
+        "conformation",
+        "neighbor_metadata",
+        "neighbor_initialise",
+        "frame",
+        "neighbor_finalise",
+    ]