begin the process of wiring the DAG for both entropy and levels

Author: harryswift01

CodeEntropy/config/run.py

@@ -17,11 +17,11 @@
 from CodeEntropy.config.arg_config_manager import ConfigManager
 from CodeEntropy.config.data_logger import DataLogger
 from CodeEntropy.config.logging_config import LoggingConfig
-from CodeEntropy.dihedral_tools import DihedralAnalysis
-from CodeEntropy.entropy import EntropyManager
-from CodeEntropy.group_molecules import GroupMolecules
-from CodeEntropy.levels import LevelManager
-from CodeEntropy.mda_universe_operations import UniverseOperations
+from CodeEntropy.entropy.entropy_manager import EntropyManager
+from CodeEntropy.group_molecules.group_molecules import GroupMolecules
+from CodeEntropy.levels.dihedral_tools import DihedralAnalysis
+from CodeEntropy.levels.level_manager import LevelManager
+from CodeEntropy.levels.mda_universe_operations import UniverseOperations
 
 logger = logging.getLogger(__name__)
 console = LoggingConfig.get_console()

CodeEntropy/entropy/entropy_graph.py

@@ -1,9 +1,9 @@
 import networkx as nx
 
-from CodeEntropy.entropy.nodes.configurational_entropy import ConfigurationalEntropyNode
-from CodeEntropy.entropy.nodes.entropy_aggregator import EntropyAggregatorNode
-from CodeEntropy.entropy.nodes.orientational_entropy import OrientationalEntropyNode
-from CodeEntropy.entropy.nodes.vibrational_entropy import VibrationalEntropyNode
+from CodeEntropy.entropy.nodes.configurational_entropy_node import (
+    ConfigurationalEntropyNode,
+)
+from CodeEntropy.entropy.nodes.vibrational_entropy_node import VibrationalEntropyNode
 
 
 class EntropyGraph:
@@ -23,23 +23,23 @@ def __init__(self):
     def build(self):
         self.add("vibrational_entropy", VibrationalEntropyNode())
 
-        self.add(
-            "orientational_entropy",
-            OrientationalEntropyNode(),
-            depends_on=["vibrational_entropy"],
-        )
+        # self.add(
+        #     "orientational_entropy",
+        #     OrientationalEntropyNode(),
+        #     depends_on=["vibrational_entropy"],
+        # )
 
         self.add(
             "configurational_entropy",
             ConfigurationalEntropyNode(),
             depends_on=["orientational_entropy"],
         )
 
-        self.add(
-            "aggregate_entropy",
-            EntropyAggregatorNode(),
-            depends_on=["configurational_entropy"],
-        )
+        # self.add(
+        #     "aggregate_entropy",
+        #     EntropyAggregatorNode(),
+        #     depends_on=["configurational_entropy"],
+        # )
 
         return self
 

CodeEntropy/entropy/entropy_manager.py

@@ -7,6 +7,7 @@
 from CodeEntropy.config.logging_config import LoggingConfig
 from CodeEntropy.entropy.entropy_graph import EntropyGraph
 from CodeEntropy.levels.hierarchy_graph import LevelDAG
+from CodeEntropy.levels.level_hierarchy import LevelHierarchy
 
 logger = logging.getLogger(__name__)
 console = LoggingConfig.get_console()
@@ -71,7 +72,9 @@ def execute(self):
 
         reduced_universe = self._get_reduced_universe()
 
-        number_molecules, levels = self._level_manager.select_levels(reduced_universe)
+        level_hierarchy = LevelHierarchy()
+
+        number_molecules, levels = level_hierarchy.select_levels(reduced_universe)
 
         groups = self._group_molecules.grouping_molecules(
             reduced_universe, self._args.grouping
@@ -112,7 +115,9 @@ def execute(self):
             "run_manager": self._run_manager,
         }
 
-        level_results = LevelDAG().build().execute(shared_data)
+        logger.info(f"shared_data: {shared_data}")
+
+        level_results = LevelDAG(self._universe_operations).build().execute(shared_data)
 
         shared_data.update(level_results)
 

CodeEntropy/entropy/nodes/configurational_entropy_node.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from CodeEntropy.entropy.configurational_entropy import ConfigurationalEntropy
+from CodeEntropy.entropy.configurational_entropy import ConformationalEntropy
 
 logger = logging.getLogger(__name__)
 
@@ -14,7 +14,7 @@ class ConfigurationalEntropyNode:
     """
 
     def __init__(self, data_logger):
-        self._ce = ConfigurationalEntropy()
+        self._ce = ConformationalEntropy()
         self._data_logger = data_logger
 
     def run(self, shared_data, **_):

CodeEntropy/levels/hierarchy_graph.py

@@ -1,44 +1,65 @@
+import logging
+
 import networkx as nx
 
 from CodeEntropy.levels.nodes.build_beads import BuildBeadsNode
-from CodeEntropy.levels.nodes.build_conformations import BuildConformationsNode
 from CodeEntropy.levels.nodes.build_covariance_matrices import (
     BuildCovarianceMatricesNode,
 )
 from CodeEntropy.levels.nodes.compute_axes import ComputeAxesNode
-from CodeEntropy.levels.nodes.compute_dihedrals import ComputeDihedralsNode
-from CodeEntropy.levels.nodes.compute_neighbours import ComputeNeighboursNode
+from CodeEntropy.levels.nodes.compute_dihedrals import ComputeConformationalStatesNode
 from CodeEntropy.levels.nodes.compute_weighted_forces import ComputeWeightedForcesNode
 from CodeEntropy.levels.nodes.compute_weighted_torques import ComputeWeightedTorquesNode
 from CodeEntropy.levels.nodes.detect_levels import DetectLevelsNode
 from CodeEntropy.levels.nodes.detect_molecules import DetectMoleculesNode
 
+logger = logging.getLogger(__name__)
+
 
 class LevelDAG:
+    """
+    DAG for computing level-resolved structural quantities.
+    Uses shared_data as the single state container.
+    """
 
-    def __init__(self):
+    def __init__(self, universe_operations):
         self.graph = nx.DiGraph()
         self.nodes = {}
+        self._universe_operations = universe_operations
 
     def build(self):
         self.add("detect_molecules", DetectMoleculesNode())
         self.add("detect_levels", DetectLevelsNode(), ["detect_molecules"])
+
         self.add("build_beads", BuildBeadsNode(), ["detect_levels"])
+
         self.add("compute_axes", ComputeAxesNode(), ["build_beads"])
+
         self.add(
-            "compute_weighted_forces", ComputeWeightedForcesNode(), ["compute_axes"]
+            "compute_weighted_forces",
+            ComputeWeightedForcesNode(),
+            ["compute_axes"],
         )
         self.add(
-            "compute_weighted_torques", ComputeWeightedTorquesNode(), ["compute_axes"]
+            "compute_weighted_torques",
+            ComputeWeightedTorquesNode(),
+            ["compute_axes"],
         )
+
         self.add(
             "build_covariance",
             BuildCovarianceMatricesNode(),
-            ["compute_weighted_forces", "compute_weighted_torques"],
+            [
+                "compute_weighted_forces",
+                "compute_weighted_torques",
+            ],
+        )
+
+        self.add(
+            "compute_conformational_states",
+            ComputeConformationalStatesNode(self._universe_operations),
+            ["detect_levels"],
         )
-        self.add("compute_dihedrals", ComputeDihedralsNode(), ["build_beads"])
-        self.add("build_conformations", BuildConformationsNode(), ["compute_dihedrals"])
-        self.add("compute_neighbours", ComputeNeighboursNode(), ["build_beads"])
 
         return self
 
@@ -50,5 +71,9 @@ def add(self, name, obj, deps=None):
                 self.graph.add_edge(d, name)
 
     def execute(self, shared_data):
-        for node in nx.topological_sort(self.graph):
-            self.nodes[node].run(shared_data)
+        """
+        Execute DAG in topological order.
+        Nodes mutate shared_data in-place.
+        """
+        for node_name in nx.topological_sort(self.graph):
+            self.nodes[node_name].run(shared_data)

CodeEntropy/levels/nodes/build_conformations.py

@@ -1,9 +1,9 @@
-from CodeEntropy.levels.dihedral_tools import DihedralTools
+from CodeEntropy.levels.dihedral_tools import DihedralAnalysis
 
 
 class BuildConformationsNode:
     def __init__(self):
-        self._dih = DihedralTools()
+        self._dih = DihedralAnalysis()
 
     def run(self, shared_data):
         dihedrals = shared_data["dihedrals"]

CodeEntropy/levels/nodes/build_covariance_matrices.py

@@ -6,18 +6,13 @@ def __init__(self):
         self._ft = ForceTorqueManager()
 
     def run(self, shared_data):
-        """
-        Build force and torque covariance matrices from weighted forces/torques.
-        """
-
-        weighted_forces = shared_data["weighted_forces"]
-        weighted_torques = shared_data["weighted_torques"]
+        forces = shared_data["weighted_forces"]
+        torques = shared_data["weighted_torques"]
 
         force_cov, torque_cov, frame_counts = self._ft.build_covariance_matrices(
-            weighted_forces,
-            weighted_torques,
+            forces, torques
         )
 
-        shared_data["force_covariance"] = force_cov
-        shared_data["torque_covariance"] = torque_cov
+        shared_data["force_covariances"] = force_cov
+        shared_data["torque_covariances"] = torque_cov
         shared_data["frame_counts"] = frame_counts

CodeEntropy/levels/nodes/compute_axes.py

@@ -1,5 +1,9 @@
+import logging
+
 from CodeEntropy.levels.coordinate_system import CoordinateSystem
 
+logger = logging.getLogger(__name__)
+
 
 class ComputeAxesNode:
     def __init__(self):
@@ -8,16 +12,26 @@ def __init__(self):
     def run(self, shared_data):
         beads = shared_data["beads"]
 
-        axes = {}
+        trans_axes = {}
+        rot_axes = {}
         avg_pos = {}
 
         for key, bead_list in beads.items():
-            axes[key] = []
+            trans_axes[key] = []
+            rot_axes[key] = []
             avg_pos[key] = []
 
             for bead in bead_list:
-                avg_pos[key].append(self._coord.get_avg_pos(bead))
-                axes[key].append(self._coord.get_axes(bead))
+                t_ax, r_ax = self._coord.get_axes(
+                    bead.data_container, bead.level, bead.index
+                )
+
+                trans_axes[key].append(t_ax)
+                rot_axes[key].append(r_ax)
+                avg_pos[key].append(
+                    self._coord.get_avg_pos(bead.atoms, bead.atoms.center_of_mass())
+                )
 
-        shared_data["axes"] = axes
+        shared_data["trans_axes"] = trans_axes
+        shared_data["rot_axes"] = rot_axes
         shared_data["avg_positions"] = avg_pos

CodeEntropy/levels/nodes/compute_dihedrals.py

@@ -1,10 +1,34 @@
-from CodeEntropy.levels.dihedral_tools import DihedralTools
+from CodeEntropy.levels.dihedral_tools import DihedralAnalysis
 
 
-class ComputeDihedralsNode:
-    def __init__(self):
-        self._dih = DihedralTools()
+class ComputeConformationalStatesNode:
+    def __init__(self, universe_operations):
+        self._dih = DihedralAnalysis(universe_operations)
 
     def run(self, shared_data):
-        beads = shared_data["beads"]
-        shared_data["dihedrals"] = self._dih.get_dihedrals(beads)
+        u = shared_data["universe"]
+        levels = shared_data["levels"]
+        groups = shared_data["groups"]
+
+        start = shared_data["start"]
+        end = shared_data["end"]
+        step = shared_data["step"]
+        bin_width = shared_data["args"].bin_width
+
+        states_ua, states_res = self._dih.build_conformational_states(
+            data_container=u,
+            levels=levels,
+            groups=groups,
+            start=start,
+            end=end,
+            step=step,
+            bin_width=bin_width,
+        )
+
+        shared_data["states_united_atom"] = states_ua
+        shared_data["states_residue"] = states_res
+
+        return {
+            "states_united_atom": states_ua,
+            "states_residue": states_res,
+        }

CodeEntropy/levels/nodes/compute_neighbours.py

@@ -1,10 +1,10 @@
-from CodeEntropy.levels.dihedral_tools import DihedralTools
+from CodeEntropy.levels.neighbours import Neighbours
 
 
-class BuildConformationsNode:
+class ComputeNeighboursNode:
     def __init__(self):
-        self._dih = DihedralTools()
+        self._nb = Neighbours()
 
     def run(self, shared_data):
-        dihedrals = shared_data["dihedrals"]
-        shared_data["conformations"] = self._dih.build_conformational_states(dihedrals)
+        beads = shared_data["beads"]
+        shared_data["neighbours"] = self._nb.compute(beads)