From 69cb883675f30a5d3de9c9ed324cc649a6272b65 Mon Sep 17 00:00:00 2001 From: harryswift01 Date: Tue, 16 Jun 2026 14:07:01 +0100 Subject: [PATCH 1/3] refactor(levels): introduce dedicated ConformationDAG stage --- .../conformations.py => conformation_dag.py} | 46 +++++------ CodeEntropy/levels/level_dag.py | 32 +++++--- ...tions_node.py => test_conformation_dag.py} | 34 +++++--- .../unit/CodeEntropy/levels/test_level_dag.py | 79 ++++++++++++++++++- 4 files changed, 136 insertions(+), 55 deletions(-) rename CodeEntropy/levels/{nodes/conformations.py => conformation_dag.py} (62%) rename tests/unit/CodeEntropy/levels/{nodes/test_conformations_node.py => test_conformation_dag.py} (77%) diff --git a/CodeEntropy/levels/nodes/conformations.py b/CodeEntropy/levels/conformation_dag.py similarity index 62% rename from CodeEntropy/levels/nodes/conformations.py rename to CodeEntropy/levels/conformation_dag.py index 5b4220a3..eb18ff96 100644 --- a/CodeEntropy/levels/nodes/conformations.py +++ b/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"] diff --git a/CodeEntropy/levels/level_dag.py b/CodeEntropy/levels/level_dag.py index ecbf1d49..cc8129a1 100644 --- a/CodeEntropy/levels/level_dag.py +++ b/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,6 +12,7 @@ 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 @@ -19,7 +20,6 @@ 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], diff --git a/tests/unit/CodeEntropy/levels/nodes/test_conformations_node.py b/tests/unit/CodeEntropy/levels/test_conformation_dag.py similarity index 77% rename from tests/unit/CodeEntropy/levels/nodes/test_conformations_node.py rename to tests/unit/CodeEntropy/levels/test_conformation_dag.py index 22b81070..ad469ffe 100644 --- a/tests/unit/CodeEntropy/levels/nodes/test_conformations_node.py +++ b/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 diff --git a/tests/unit/CodeEntropy/levels/test_level_dag.py b/tests/unit/CodeEntropy/levels/test_level_dag.py index 30ae4077..1b9f4d4f 100644 --- a/tests/unit/CodeEntropy/levels/test_level_dag.py +++ b/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", + ] From b1cc6995f41260eca2a4a826df07dcced6d5336e Mon Sep 17 00:00:00 2001 From: harryswift01 Date: Tue, 16 Jun 2026 14:57:21 +0100 Subject: [PATCH 2/3] docs(api): update ConformationDAG autodocs --- docs/api/CodeEntropy.levels.conformation_dag.rst | 7 +++++++ docs/api/CodeEntropy.levels.nodes.conformations.rst | 7 ------- docs/api/CodeEntropy.levels.nodes.rst | 1 - docs/api/CodeEntropy.levels.rst | 1 + 4 files changed, 8 insertions(+), 8 deletions(-) create mode 100644 docs/api/CodeEntropy.levels.conformation_dag.rst delete mode 100644 docs/api/CodeEntropy.levels.nodes.conformations.rst diff --git a/docs/api/CodeEntropy.levels.conformation_dag.rst b/docs/api/CodeEntropy.levels.conformation_dag.rst new file mode 100644 index 00000000..ffea54eb --- /dev/null +++ b/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: diff --git a/docs/api/CodeEntropy.levels.nodes.conformations.rst b/docs/api/CodeEntropy.levels.nodes.conformations.rst deleted file mode 100644 index 7bd29ddb..00000000 --- a/docs/api/CodeEntropy.levels.nodes.conformations.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.nodes.conformations module -============================================= - -.. automodule:: CodeEntropy.levels.nodes.conformations - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.nodes.rst b/docs/api/CodeEntropy.levels.nodes.rst index 7ce5004f..4499f939 100644 --- a/docs/api/CodeEntropy.levels.nodes.rst +++ b/docs/api/CodeEntropy.levels.nodes.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 diff --git a/docs/api/CodeEntropy.levels.rst b/docs/api/CodeEntropy.levels.rst index a521eb6b..cca6398d 100644 --- a/docs/api/CodeEntropy.levels.rst +++ b/docs/api/CodeEntropy.levels.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 From 135428ba7ded0ba87407c2c7696f5f7988f81552 Mon Sep 17 00:00:00 2001 From: harryswift01 Date: Tue, 16 Jun 2026 15:03:46 +0100 Subject: [PATCH 3/3] docs(doc-string): update doc-strings for ConformationalDAG --- CodeEntropy/levels/conformation_dag.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/CodeEntropy/levels/conformation_dag.py b/CodeEntropy/levels/conformation_dag.py index eb18ff96..78ad12ab 100644 --- a/CodeEntropy/levels/conformation_dag.py +++ b/CodeEntropy/levels/conformation_dag.py @@ -17,12 +17,7 @@ class ConformationDAG: - """Execute conformational-state construction for selected trajectory frames. - - The first implementation intentionally preserves the existing serial - ConformationStateBuilder behaviour. Later issues can replace this internal - implementation with chunked map-reduce execution. - """ + """Execute conformational-state construction for selected trajectory frames.""" def __init__(self, universe_operations: Any | None = None) -> None: self._builder = ConformationStateBuilder( @@ -43,7 +38,15 @@ def execute( *, progress: object | None = None, ) -> dict[str, ConformationalStates]: - """Compute conformational states and store them in shared workflow data.""" + """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. + """ universe = shared_data["reduced_universe"] levels = shared_data["levels"] groups = shared_data["groups"]