diff --git a/.github/workflows/daily.yaml b/.github/workflows/daily.yaml index 79a5b4b..167ea27 100644 --- a/.github/workflows/daily.yaml +++ b/.github/workflows/daily.yaml @@ -9,6 +9,13 @@ concurrency: group: daily-${{ github.ref }} cancel-in-progress: true +env: + OMP_NUM_THREADS: 1 + MKL_NUM_THREADS: 1 + OPENBLAS_NUM_THREADS: 1 + NUMEXPR_NUM_THREADS: 1 + PYTHONHASHSEED: 0 + jobs: unit: name: Unit (${{ matrix.os }}, py${{ matrix.python-version }}) @@ -37,4 +44,4 @@ jobs: python -m pip install -e .[testing] - name: Run unit tests - run: python -m pytest tests/unit + run: python -m pytest tests/unit -n auto --dist=loadscope diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 96aaa0c..285c504 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -7,11 +7,19 @@ concurrency: group: pr-${{ github.ref }} cancel-in-progress: true +env: + OMP_NUM_THREADS: 1 + MKL_NUM_THREADS: 1 + OPENBLAS_NUM_THREADS: 1 + NUMEXPR_NUM_THREADS: 1 + PYTHONHASHSEED: 0 + jobs: unit: name: Unit runs-on: ${{ matrix.os }} timeout-minutes: 25 + strategy: fail-fast: false matrix: @@ -33,8 +41,10 @@ jobs: python -m pip install --upgrade pip python -m pip install -e .[testing] - - name: Pytest (unit) • ${{ matrix.os }}, ${{ matrix.python-version }} - run: python -m pytest tests/unit + - name: Run unit tests + run: | + python -m pytest tests/unit -n auto --dist=loadscope + discover-systems: name: Discover regression systems @@ -63,6 +73,7 @@ jobs: SYSTEMS=$(python -m tests.regression.list_systems) echo "systems=$SYSTEMS" >> $GITHUB_OUTPUT + regression-quick: name: Regression (fast) • ${{ matrix.system }} needs: [unit, discover-systems] @@ -84,23 +95,16 @@ jobs: python-version: "3.14" cache: pip - - name: Cache testdata - uses: actions/cache@v4 - with: - path: .testdata - key: codeentropy-testdata-${{ runner.os }}-py314 - - name: Install testing dependencies run: | python -m pip install --upgrade pip python -m pip install -e .[testing] - - name: Run fast regression tests (per system) + - name: Run fast regression tests run: | python -m pytest tests/regression \ -m "not slow" \ - -n auto \ - --dist=loadscope \ + -n 0 \ -k "${{ matrix.system }}" \ -vv \ --durations=20 @@ -114,6 +118,7 @@ jobs: .testdata/** /tmp/pytest-of-*/pytest-*/** + docs: name: Docs needs: unit @@ -147,8 +152,10 @@ jobs: name: docs-html path: docs/_build/html + pre-commit: name: Pre-commit + needs: unit runs-on: ubuntu-24.04 timeout-minutes: 15 @@ -176,6 +183,7 @@ jobs: exit 1 } + coverage: name: Coverage needs: unit @@ -203,6 +211,8 @@ jobs: --cov CodeEntropy \ --cov-report term-missing \ --cov-report xml \ + -n auto \ + --dist=loadscope \ -q - name: Upload coverage to Coveralls diff --git a/.github/workflows/weekly-regression.yaml b/.github/workflows/weekly-regression.yaml index 5f4595b..fb1dc29 100644 --- a/.github/workflows/weekly-regression.yaml +++ b/.github/workflows/weekly-regression.yaml @@ -9,6 +9,13 @@ concurrency: group: weekly-regression-${{ github.ref }} cancel-in-progress: true +env: + OMP_NUM_THREADS: 1 + MKL_NUM_THREADS: 1 + OPENBLAS_NUM_THREADS: 1 + NUMEXPR_NUM_THREADS: 1 + PYTHONHASHSEED: 0 + jobs: discover: name: Discover regression systems @@ -75,7 +82,7 @@ jobs: pytest tests/regression \ -k "${{ matrix.system }}" \ --run-slow \ - -n auto \ + -n 1 \ --dist=loadscope \ -vv \ --durations=20 diff --git a/CodeEntropy/entropy/workflow.py b/CodeEntropy/entropy/workflow.py index b35e259..bf770bf 100644 --- a/CodeEntropy/entropy/workflow.py +++ b/CodeEntropy/entropy/workflow.py @@ -165,7 +165,7 @@ def _build_shared_data( "universe": self._universe, "reduced_universe": reduced_universe, "levels": levels, - "groups": dict(groups), + "groups": dict(sorted(groups.items())), "start": traj.start, "end": traj.end, "step": traj.step, @@ -319,7 +319,7 @@ def _compute_water_entropy( water_entropy = WaterEntropy(self._args, self._reporter) - for group_id in water_groups.keys(): + for group_id in sorted(water_groups.keys()): water_entropy.calculate_and_log( universe=self._universe, start=traj.start, diff --git a/CodeEntropy/levels/dihedrals.py b/CodeEntropy/levels/dihedrals.py index 0b40278..96df364 100644 --- a/CodeEntropy/levels/dihedrals.py +++ b/CodeEntropy/levels/dihedrals.py @@ -105,7 +105,7 @@ def build_conformational_states( progress.advance(task) return states_ua, states_res - for group_id in groups.keys(): + for group_id in sorted(groups.keys()): molecules = groups[group_id] if not molecules: if progress is not None and task is not None: diff --git a/CodeEntropy/levels/level_dag.py b/CodeEntropy/levels/level_dag.py index f43f1ef..874fe17 100644 --- a/CodeEntropy/levels/level_dag.py +++ b/CodeEntropy/levels/level_dag.py @@ -11,7 +11,7 @@ 2) Frame stage (runs for each trajectory frame): - Execute the `FrameGraph` to produce frame-local covariance outputs. - - Reduce frame-local outputs into running (incremental) means. + - Reduce frame-local outputs into deterministic sums and counts. """ from __future__ import annotations @@ -41,10 +41,11 @@ class LevelDAG: The LevelDAG is responsible for: - Running a static DAG (once) to prepare shared inputs. - Running a per-frame DAG (for each frame) to compute frame-local outputs. - - Reducing frame-local outputs into shared running means. + - Reducing frame-local outputs into deterministic sums and counts. - The reduction performed here is an incremental mean across frames (and across - molecules within a group when frame nodes average within-frame first). + The reduction performed here is order-independent: frame-local sums and + counts are accumulated across frames and final means are computed once after + all frames have been processed. """ def __init__(self, universe_operations: Any | None = None) -> None: @@ -98,7 +99,7 @@ def execute( This method ensures required shared components exist, runs the static stage once, then iterates through trajectory frames to run the per-frame stage and - reduce outputs into running means. + reduce outputs into deterministic sums and counts. Args: shared_data: Shared workflow data dict. This mapping is mutated in-place @@ -112,6 +113,7 @@ def execute( shared_data.setdefault("axes_manager", AxesCalculator()) self._run_static_stage(shared_data, progress=progress) self._run_frame_stage(shared_data, progress=progress) + self._finalize_means(shared_data) return shared_data def _run_static_stage( @@ -220,26 +222,10 @@ def _run_frame_stage( if progress is not None and task is not None: progress.advance(task) - @staticmethod - def _incremental_mean(old: Any, new: Any, n: int) -> Any: - """Compute an incremental mean. - - Args: - old: Previous running mean (or None for first sample). - new: New sample to incorporate. - n: 1-based sample count after adding `new`. - - Returns: - Updated running mean. - """ - if old is None: - return new.copy() if hasattr(new, "copy") else new - return old + (new - old) / float(n) - def _reduce_one_frame( self, shared_data: dict[str, Any], frame_out: dict[str, Any] ) -> None: - """Reduce one frame's covariance outputs into shared running means. + """Reduce one frame's covariance outputs into shared sum accumulators. Args: shared_data: Shared workflow data dict containing accumulators. @@ -251,94 +237,191 @@ def _reduce_one_frame( def _reduce_force_and_torque( self, shared_data: dict[str, Any], frame_out: dict[str, Any] ) -> None: - """Reduce force/torque covariance outputs into shared accumulators. + """Reduce force/torque frame-local sums into shared accumulators. Args: shared_data: Shared workflow data dict containing: - - "force_covariances", "torque_covariances": accumulator structures. - - "frame_counts": running sample counts for each accumulator slot. + - "force_sums", "torque_sums": running sum accumulators. + - "force_counts", "torque_counts": running sample counts. - "group_id_to_index": mapping from group id to accumulator index. - frame_out: Frame-local outputs containing "force" and "torque" sections. + frame_out: Frame-local outputs containing "force", "torque", + "force_counts", and "torque_counts" sections. Returns: - None. Mutates accumulator values and counts in shared_data in-place. + None. Mutates shared accumulators and counts in-place. """ - f_cov = shared_data["force_covariances"] - t_cov = shared_data["torque_covariances"] - counts = shared_data["frame_counts"] + f_sums = shared_data["force_sums"] + t_sums = shared_data["torque_sums"] + f_counts = shared_data["force_counts"] + t_counts = shared_data["torque_counts"] gid2i = shared_data["group_id_to_index"] f_frame = frame_out["force"] t_frame = frame_out["torque"] - - for key, F in f_frame["ua"].items(): - counts["ua"][key] = counts["ua"].get(key, 0) + 1 - n = counts["ua"][key] - f_cov["ua"][key] = self._incremental_mean(f_cov["ua"].get(key), F, n) - - for key, T in t_frame["ua"].items(): - if key not in counts["ua"]: - counts["ua"][key] = counts["ua"].get(key, 0) + 1 - n = counts["ua"][key] - t_cov["ua"][key] = self._incremental_mean(t_cov["ua"].get(key), T, n) - - for gid, F in f_frame["res"].items(): + f_frame_counts = frame_out["force_counts"] + t_frame_counts = frame_out["torque_counts"] + + for key in sorted(f_frame["ua"].keys()): + F = f_frame["ua"][key] + c = int(f_frame_counts["ua"].get(key, 0)) + if c <= 0: + continue + prev = f_sums["ua"].get(key) + f_sums["ua"][key] = F.copy() if prev is None else prev + F + f_counts["ua"][key] = f_counts["ua"].get(key, 0) + c + + for key in sorted(t_frame["ua"].keys()): + T = t_frame["ua"][key] + c = int(t_frame_counts["ua"].get(key, 0)) + if c <= 0: + continue + prev = t_sums["ua"].get(key) + t_sums["ua"][key] = T.copy() if prev is None else prev + T + t_counts["ua"][key] = t_counts["ua"].get(key, 0) + c + + for gid in sorted(f_frame["res"].keys()): + F = f_frame["res"][gid] gi = gid2i[gid] - counts["res"][gi] += 1 - n = counts["res"][gi] - f_cov["res"][gi] = self._incremental_mean(f_cov["res"][gi], F, n) - - for gid, T in t_frame["res"].items(): + c = int(f_frame_counts["res"].get(gid, 0)) + if c <= 0: + continue + prev = f_sums["res"][gi] + f_sums["res"][gi] = F.copy() if prev is None else prev + F + f_counts["res"][gi] += c + + for gid in sorted(t_frame["res"].keys()): + T = t_frame["res"][gid] gi = gid2i[gid] - if counts["res"][gi] == 0: - counts["res"][gi] += 1 - n = counts["res"][gi] - t_cov["res"][gi] = self._incremental_mean(t_cov["res"][gi], T, n) - - for gid, F in f_frame["poly"].items(): + c = int(t_frame_counts["res"].get(gid, 0)) + if c <= 0: + continue + prev = t_sums["res"][gi] + t_sums["res"][gi] = T.copy() if prev is None else prev + T + t_counts["res"][gi] += c + + for gid in sorted(f_frame["poly"].keys()): + F = f_frame["poly"][gid] gi = gid2i[gid] - counts["poly"][gi] += 1 - n = counts["poly"][gi] - f_cov["poly"][gi] = self._incremental_mean(f_cov["poly"][gi], F, n) - - for gid, T in t_frame["poly"].items(): + c = int(f_frame_counts["poly"].get(gid, 0)) + if c <= 0: + continue + prev = f_sums["poly"][gi] + f_sums["poly"][gi] = F.copy() if prev is None else prev + F + f_counts["poly"][gi] += c + + for gid in sorted(t_frame["poly"].keys()): + T = t_frame["poly"][gid] gi = gid2i[gid] - if counts["poly"][gi] == 0: - counts["poly"][gi] += 1 - n = counts["poly"][gi] - t_cov["poly"][gi] = self._incremental_mean(t_cov["poly"][gi], T, n) + c = int(t_frame_counts["poly"].get(gid, 0)) + if c <= 0: + continue + prev = t_sums["poly"][gi] + t_sums["poly"][gi] = T.copy() if prev is None else prev + T + t_counts["poly"][gi] += c def _reduce_forcetorque( self, shared_data: dict[str, Any], frame_out: dict[str, Any] ) -> None: - """Reduce combined force-torque covariance outputs into shared accumulators. + """Reduce combined force-torque frame-local sums into shared accumulators. Args: shared_data: Shared workflow data dict containing: - - "forcetorque_covariances": accumulator structures. - - "forcetorque_counts": running sample counts for each accumulator slot. + - "forcetorque_sums": running sum accumulators. + - "forcetorque_counts": running sample counts. - "group_id_to_index": mapping from group id to accumulator index. - frame_out: Frame-local outputs that may include a "forcetorque" section. + frame_out: Frame-local outputs that may include "forcetorque" and + "forcetorque_counts" sections. Returns: - None. Mutates accumulator values and counts in shared_data in-place. + None. Mutates shared accumulators and counts in-place. """ if "forcetorque" not in frame_out: return - ft_cov = shared_data["forcetorque_covariances"] + ft_sums = shared_data["forcetorque_sums"] ft_counts = shared_data["forcetorque_counts"] gid2i = shared_data["group_id_to_index"] + ft_frame = frame_out["forcetorque"] + ft_frame_counts = frame_out.get("forcetorque_counts", {"res": {}, "poly": {}}) - for gid, M in ft_frame.get("res", {}).items(): + for gid in sorted(ft_frame.get("res", {}).keys()): + M = ft_frame["res"][gid] gi = gid2i[gid] - ft_counts["res"][gi] += 1 - n = ft_counts["res"][gi] - ft_cov["res"][gi] = self._incremental_mean(ft_cov["res"][gi], M, n) - - for gid, M in ft_frame.get("poly", {}).items(): + c = int(ft_frame_counts.get("res", {}).get(gid, 0)) + if c <= 0: + continue + prev = ft_sums["res"][gi] + ft_sums["res"][gi] = M.copy() if prev is None else prev + M + ft_counts["res"][gi] += c + + for gid in sorted(ft_frame.get("poly", {}).keys()): + M = ft_frame["poly"][gid] gi = gid2i[gid] - ft_counts["poly"][gi] += 1 - n = ft_counts["poly"][gi] - ft_cov["poly"][gi] = self._incremental_mean(ft_cov["poly"][gi], M, n) + c = int(ft_frame_counts.get("poly", {}).get(gid, 0)) + if c <= 0: + continue + prev = ft_sums["poly"][gi] + ft_sums["poly"][gi] = M.copy() if prev is None else prev + M + ft_counts["poly"][gi] += c + + def _finalize_means(self, shared_data: dict[str, Any]) -> None: + """Compute finalized mean matrices from accumulated sums and counts. + + Args: + shared_data: Shared workflow data dict containing running sums and counts. + + Returns: + None. Writes finalized mean matrices back into shared_data. + """ + + def _compute_means( + sums: dict[str, Any], + counts: dict[str, Any], + ) -> dict[str, Any]: + out: dict[str, Any] = {} + + for domain in sorted(sums.keys()): + domain_sums = sums[domain] + domain_counts = counts[domain] + + if isinstance(domain_sums, dict): + out[domain] = {} + for key in sorted(domain_sums.keys()): + total = domain_sums[key] + count = int(domain_counts.get(key, 0)) + out[domain][key] = total / float(count) if count > 0 else None + continue + + mean_list: list[Any] = [None] * len(domain_sums) + for idx, total in enumerate(domain_sums): + if total is None: + continue + count = int(domain_counts[idx]) + mean_list[idx] = total / float(count) if count > 0 else None + out[domain] = mean_list + + return out + + shared_data["force_covariances"] = _compute_means( + shared_data["force_sums"], + shared_data["force_counts"], + ) + shared_data["torque_covariances"] = _compute_means( + shared_data["torque_sums"], + shared_data["torque_counts"], + ) + shared_data["forcetorque_covariances"] = _compute_means( + shared_data["forcetorque_sums"], + shared_data["forcetorque_counts"], + ) + + shared_data["frame_counts"] = shared_data["force_counts"] + shared_data["force_torque_stats"] = { + "res": list(shared_data["forcetorque_covariances"]["res"]), + "poly": list(shared_data["forcetorque_covariances"]["poly"]), + } + shared_data["force_torque_counts"] = { + "res": shared_data["forcetorque_counts"]["res"].copy(), + "poly": shared_data["forcetorque_counts"]["poly"].copy(), + } diff --git a/CodeEntropy/levels/neighbors.py b/CodeEntropy/levels/neighbors.py index a91b395..5a07531 100644 --- a/CodeEntropy/levels/neighbors.py +++ b/CodeEntropy/levels/neighbors.py @@ -56,7 +56,7 @@ def get_neighbors(self, universe, levels, groups, search_type): number_frames = len(universe.trajectory) - for group_id in groups.keys(): + for group_id in sorted(groups.keys()): molecules = groups[group_id] highest_level = levels[molecules[0]][-1] @@ -120,7 +120,7 @@ def get_symmetry(self, universe, groups): symmetry_number = {} linear = {} - for group_id in groups.keys(): + for group_id in sorted(groups.keys()): molecules = groups[group_id] rdkit_mol, number_heavy, number_hydrogen = self._get_rdkit_mol( diff --git a/CodeEntropy/levels/nodes/accumulators.py b/CodeEntropy/levels/nodes/accumulators.py index 722cc96..2c6f1a6 100644 --- a/CodeEntropy/levels/nodes/accumulators.py +++ b/CodeEntropy/levels/nodes/accumulators.py @@ -1,12 +1,12 @@ """Initialize covariance accumulators. This module defines a LevelDAG static node that allocates all per-frame reduction -accumulators (means) and counters used by downstream frame processing. +accumulators (sums) and counters used by downstream frame processing. The node owns only initialization concerns (single responsibility): - create group-id <-> index mappings -- allocate force/torque covariance mean containers -- allocate optional combined force-torque (FT) mean containers +- allocate force/torque covariance sum containers +- allocate optional combined force-torque (FT) sum containers - allocate per-level frame counters The structure created here is treated as the canonical storage layout for the @@ -37,12 +37,13 @@ class GroupIndex: @dataclass(frozen=True) class CovarianceAccumulators: - """Container for covariance mean accumulators and frame counters.""" + """Container for covariance sum accumulators and frame counters.""" - force_covariances: dict[str, Any] - torque_covariances: dict[str, Any] - frame_counts: dict[str, Any] - forcetorque_covariances: dict[str, Any] + force_sums: dict[str, Any] + torque_sums: dict[str, Any] + force_counts: dict[str, Any] + torque_counts: dict[str, Any] + forcetorque_sums: dict[str, Any] forcetorque_counts: dict[str, Any] @@ -51,21 +52,29 @@ class InitCovarianceAccumulatorsNode: Produces the following keys in `shared_data`: - Canonical mean accumulators: - - force_covariances: {"ua": dict, "res": list, "poly": list} - - torque_covariances: {"ua": dict, "res": list, "poly": list} - - forcetorque_covariances: {"res": list, "poly": list} (6N x 6N means) + Canonical sum accumulators: + - force_sums: {"ua": dict, "res": list, "poly": list} + - torque_sums: {"ua": dict, "res": list, "poly": list} + - forcetorque_sums: {"res": list, "poly": list} (6N x 6N sums) Counters: - - frame_counts: {"ua": dict, "res": np.ndarray[int], "poly": np.ndarray[int]} + - force_counts: {"ua": dict, "res": np.ndarray[int], "poly": np.ndarray[int]} + - torque_counts: {"ua": dict, "res": np.ndarray[int], "poly": np.ndarray[int]} - forcetorque_counts: {"res": np.ndarray[int], "poly": np.ndarray[int]} Group index mapping: - group_id_to_index: {group_id: index} - index_to_group_id: [group_id_by_index] - Backwards-compatible aliases (kept for older consumers): - - force_torque_stats -> forcetorque_covariances + Compatibility aliases: + - force_covariances -> force_sums during reduction, later overwritten + with finalized means by LevelDAG. + - torque_covariances -> torque_sums during reduction, later overwritten + with finalized means by LevelDAG. + - forcetorque_covariances -> forcetorque_sums during reduction, later + overwritten with finalized means by LevelDAG. + - frame_counts -> force_counts + - force_torque_stats -> forcetorque_sums - force_torque_counts -> forcetorque_counts """ @@ -89,7 +98,7 @@ def run(self, shared_data: dict[str, Any]) -> dict[str, Any]: ) self._attach_to_shared_data(shared_data, group_index, accumulators) - self._attach_backwards_compatible_aliases(shared_data) + self._attach_compatible_aliases(shared_data) return self._build_return_payload(shared_data) @@ -103,13 +112,13 @@ def _build_group_index(groups: dict[int, Any]) -> GroupIndex: Returns: GroupIndex mapping object. """ - group_ids = list(groups.keys()) + group_ids = sorted(groups.keys()) gid2i = {gid: i for i, gid in enumerate(group_ids)} return GroupIndex(group_id_to_index=gid2i, index_to_group_id=list(group_ids)) @staticmethod def _build_accumulators(n_groups: int) -> CovarianceAccumulators: - """Allocate empty covariance means and counters. + """Allocate empty covariance sum containers and counters. Args: n_groups: Number of molecule groups. @@ -117,26 +126,32 @@ def _build_accumulators(n_groups: int) -> CovarianceAccumulators: Returns: CovarianceAccumulators containing allocated containers. """ - force_cov = {"ua": {}, "res": [None] * n_groups, "poly": [None] * n_groups} - torque_cov = {"ua": {}, "res": [None] * n_groups, "poly": [None] * n_groups} + force_sums = {"ua": {}, "res": [None] * n_groups, "poly": [None] * n_groups} + torque_sums = {"ua": {}, "res": [None] * n_groups, "poly": [None] * n_groups} - frame_counts = { + force_counts = { + "ua": {}, + "res": np.zeros(n_groups, dtype=int), + "poly": np.zeros(n_groups, dtype=int), + } + torque_counts = { "ua": {}, "res": np.zeros(n_groups, dtype=int), "poly": np.zeros(n_groups, dtype=int), } - forcetorque_cov = {"res": [None] * n_groups, "poly": [None] * n_groups} + forcetorque_sums = {"res": [None] * n_groups, "poly": [None] * n_groups} forcetorque_counts = { "res": np.zeros(n_groups, dtype=int), "poly": np.zeros(n_groups, dtype=int), } return CovarianceAccumulators( - force_covariances=force_cov, - torque_covariances=torque_cov, - frame_counts=frame_counts, - forcetorque_covariances=forcetorque_cov, + force_sums=force_sums, + torque_sums=torque_sums, + force_counts=force_counts, + torque_counts=torque_counts, + forcetorque_sums=forcetorque_sums, forcetorque_counts=forcetorque_counts, ) @@ -154,21 +169,27 @@ def _attach_to_shared_data( shared_data["group_id_to_index"] = group_index.group_id_to_index shared_data["index_to_group_id"] = group_index.index_to_group_id - shared_data["force_covariances"] = acc.force_covariances - shared_data["torque_covariances"] = acc.torque_covariances - shared_data["frame_counts"] = acc.frame_counts + shared_data["force_sums"] = acc.force_sums + shared_data["torque_sums"] = acc.torque_sums + shared_data["force_counts"] = acc.force_counts + shared_data["torque_counts"] = acc.torque_counts - shared_data["forcetorque_covariances"] = acc.forcetorque_covariances + shared_data["forcetorque_sums"] = acc.forcetorque_sums shared_data["forcetorque_counts"] = acc.forcetorque_counts @staticmethod - def _attach_backwards_compatible_aliases(shared_data: SharedData) -> None: - """Attach backwards-compatible aliases. + def _attach_compatible_aliases(shared_data: SharedData) -> None: + """Attach compatibility aliases. Args: shared_data: Shared pipeline dictionary. """ - shared_data["force_torque_stats"] = shared_data["forcetorque_covariances"] + shared_data["force_covariances"] = shared_data["force_sums"] + shared_data["torque_covariances"] = shared_data["torque_sums"] + shared_data["forcetorque_covariances"] = shared_data["forcetorque_sums"] + + shared_data["frame_counts"] = shared_data["force_counts"] + shared_data["force_torque_stats"] = shared_data["forcetorque_sums"] shared_data["force_torque_counts"] = shared_data["forcetorque_counts"] @staticmethod @@ -184,11 +205,16 @@ def _build_return_payload(shared_data: SharedData) -> dict[str, Any]: return { "group_id_to_index": shared_data["group_id_to_index"], "index_to_group_id": shared_data["index_to_group_id"], + "force_sums": shared_data["force_sums"], + "torque_sums": shared_data["torque_sums"], + "force_counts": shared_data["force_counts"], + "torque_counts": shared_data["torque_counts"], + "forcetorque_sums": shared_data["forcetorque_sums"], + "forcetorque_counts": shared_data["forcetorque_counts"], "force_covariances": shared_data["force_covariances"], "torque_covariances": shared_data["torque_covariances"], "frame_counts": shared_data["frame_counts"], "forcetorque_covariances": shared_data["forcetorque_covariances"], - "forcetorque_counts": shared_data["forcetorque_counts"], "force_torque_stats": shared_data["force_torque_stats"], "force_torque_counts": shared_data["force_torque_counts"], } diff --git a/CodeEntropy/levels/nodes/covariance.py b/CodeEntropy/levels/nodes/covariance.py index 25375d4..0128698 100644 --- a/CodeEntropy/levels/nodes/covariance.py +++ b/CodeEntropy/levels/nodes/covariance.py @@ -2,13 +2,14 @@ This module computes per-frame second-moment matrices for force and torque vectors at each hierarchy level (united_atom, residue, polymer). Results are -incrementally averaged across molecules within a group for the current frame. +accumulated as deterministic sums and counts across molecules within a group +for the current frame. Responsibilities: - Build bead-level force/torque vectors using ForceTorqueCalculator. - Construct per-frame force/torque second moments (outer products). - Optionally construct combined force-torque block matrices. -- Average per-frame matrices across molecules in the same group. +- Accumulate per-frame matrices and counts across molecules in the same group. Not responsible for: - Defining groups/levels/beads mapping (provided via shared context). @@ -42,9 +43,9 @@ class FrameCovarianceNode: - residue - polymer - Within a single frame, outputs are incrementally averaged across molecules - that belong to the same group. Frame-to-frame accumulation is handled - elsewhere (by a higher-level reducer). + Within a single frame, outputs are accumulated as sums together with sample + counts across molecules that belong to the same group. Frame-to-frame + accumulation is handled elsewhere (by a higher-level reducer). """ @@ -56,16 +57,16 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]: """Compute and store per-frame force/torque (and optional FT) matrices. Args: - ctx: Frame context dict expected to include: - - "shared": dict containing reduced_universe, groups, levels, beads, - args - - shared["axes_manager"] (created in static stage) + ctx: Frame context dictionary. Expected to include ``"shared"``, + containing reduced universe, groups, levels, beads, args, + and ``shared["axes_manager"]`` created during the static stage. Returns: - The frame covariance payload also stored at ctx["frame_covariance"]. + The frame covariance payload also stored at + ``ctx["frame_covariance"]``. Raises: - KeyError: If ctx is missing required fields. + KeyError: If ``ctx`` is missing required fields. """ shared = self._get_shared(ctx) @@ -85,16 +86,17 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]: out_force: dict[str, dict[Any, Matrix]] = {"ua": {}, "res": {}, "poly": {}} out_torque: dict[str, dict[Any, Matrix]] = {"ua": {}, "res": {}, "poly": {}} + out_counts: dict[str, dict[Any, int]] = {"ua": {}, "res": {}, "poly": {}} + out_ft: dict[str, dict[Any, Matrix]] | None = ( {"ua": {}, "res": {}, "poly": {}} if combined else None ) - - ua_molcount: dict[tuple[int, int], int] = {} - res_molcount: dict[int, int] = {} - poly_molcount: dict[int, int] = {} + out_ft_counts: dict[str, dict[Any, int]] | None = ( + {"ua": {}, "res": {}, "poly": {}} if combined else None + ) for group_id, mol_ids in sorted(groups.items()): - for mol_id in mol_ids: + for mol_id in sorted(mol_ids): mol = fragments[mol_id] level_list = levels[mol_id] @@ -112,7 +114,7 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]: is_highest=("united_atom" == level_list[-1]), out_force=out_force, out_torque=out_torque, - molcount=ua_molcount, + out_counts=out_counts, ) if "residue" in level_list: @@ -129,8 +131,9 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]: is_highest=("residue" == level_list[-1]), out_force=out_force, out_torque=out_torque, + out_counts=out_counts, out_ft=out_ft, - molcount=res_molcount, + out_ft_counts=out_ft_counts, combined=combined, ) @@ -147,14 +150,25 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]: is_highest=("polymer" == level_list[-1]), out_force=out_force, out_torque=out_torque, + out_counts=out_counts, out_ft=out_ft, - molcount=poly_molcount, + out_ft_counts=out_ft_counts, combined=combined, ) - frame_cov: dict[str, Any] = {"force": out_force, "torque": out_torque} - if combined and out_ft is not None: + frame_cov: dict[str, Any] = { + "force": out_force, + "torque": out_torque, + "force_counts": out_counts, + "torque_counts": { + "ua": dict(out_counts["ua"]), + "res": dict(out_counts["res"]), + "poly": dict(out_counts["poly"]), + }, + } + if combined and out_ft is not None and out_ft_counts is not None: frame_cov["forcetorque"] = out_ft + frame_cov["forcetorque_counts"] = out_ft_counts ctx["frame_covariance"] = frame_cov return frame_cov @@ -174,7 +188,7 @@ def _process_united_atom( is_highest: bool, out_force: dict[str, dict[Any, Matrix]], out_torque: dict[str, dict[Any, Matrix]], - molcount: dict[tuple[int, int], int], + out_counts: dict[str, dict[Any, int]], ) -> None: """Compute UA-level force/torque second moments for one molecule. @@ -195,10 +209,10 @@ def _process_united_atom( is_highest: Whether the UA level is the highest level for the molecule. out_force: Output accumulator for UA force second moments. out_torque: Output accumulator for UA torque second moments. - molcount: Per-(group_id, local_res_i) molecule counters for averaging. + out_counts: Output accumulator for UA molecule counts. Returns: - None. Mutates out_force/out_torque and molcount in-place. + None. Mutates out_force/out_torque and out_counts in-place. """ for local_res_i, res in enumerate(mol.residues): bead_key = (mol_id, "united_atom", local_res_i) @@ -223,10 +237,9 @@ def _process_united_atom( F, T = self._ft.compute_frame_covariance(force_vecs, torque_vecs) key = (group_id, local_res_i) - n = molcount.get(key, 0) + 1 - out_force["ua"][key] = self._inc_mean(out_force["ua"].get(key), F, n) - out_torque["ua"][key] = self._inc_mean(out_torque["ua"].get(key), T, n) - molcount[key] = n + out_force["ua"][key] = self._accumulate_sum(out_force["ua"].get(key), F) + out_torque["ua"][key] = self._accumulate_sum(out_torque["ua"].get(key), T) + out_counts["ua"][key] = out_counts["ua"].get(key, 0) + 1 def _process_residue( self, @@ -243,8 +256,9 @@ def _process_residue( is_highest: bool, out_force: dict[str, dict[Any, Matrix]], out_torque: dict[str, dict[Any, Matrix]], + out_counts: dict[str, dict[Any, int]], out_ft: dict[str, dict[Any, Matrix]] | None, - molcount: dict[int, int], + out_ft_counts: dict[str, dict[Any, int]] | None, combined: bool, ) -> None: """Compute residue-level force/torque (and optional FT) moments for one @@ -252,9 +266,9 @@ def _process_residue( Residue bead vectors are constructed for the molecule and used to compute per-frame force and torque second-moment matrices. Outputs are then - incrementally averaged across molecules in the same group for this frame. - If combined FT matrices are enabled and this is the highest level, a - force-torque block matrix is also constructed and averaged. + accumulated as sums and counts across molecules in the same group for this + frame. If combined FT matrices are enabled and this is the highest level, + a force-torque block matrix is also constructed and averaged. Args: u: MDAnalysis Universe (or compatible) providing atom access. @@ -269,12 +283,13 @@ def _process_residue( is_highest: Whether residue level is the highest level for the molecule. out_force: Output accumulator for residue force second moments. out_torque: Output accumulator for residue torque second moments. + out_counts: Output accumulator for residue molecule counts. out_ft: Optional output accumulator for residue combined FT matrices. - molcount: Per-group molecule counter for within-frame averaging. + out_ft_counts: Optional output accumulator for residue FT counts. combined: Whether combined force-torque matrices are enabled. Returns: - None. Mutates output dictionaries and molcount in-place. + None. Mutates output dictionaries and count accumulators in-place. """ bead_key = (mol_id, "residue") bead_idx_list = beads.get(bead_key, []) @@ -297,18 +312,20 @@ def _process_residue( F, T = self._ft.compute_frame_covariance(force_vecs, torque_vecs) - n = molcount.get(group_id, 0) + 1 - out_force["res"][group_id] = self._inc_mean( - out_force["res"].get(group_id), F, n + out_force["res"][group_id] = self._accumulate_sum( + out_force["res"].get(group_id), F ) - out_torque["res"][group_id] = self._inc_mean( - out_torque["res"].get(group_id), T, n + out_torque["res"][group_id] = self._accumulate_sum( + out_torque["res"].get(group_id), T ) - molcount[group_id] = n + out_counts["res"][group_id] = out_counts["res"].get(group_id, 0) + 1 - if combined and is_highest and out_ft is not None: + if combined and is_highest and out_ft is not None and out_ft_counts is not None: M = self._build_ft_block(force_vecs, torque_vecs) - out_ft["res"][group_id] = self._inc_mean(out_ft["res"].get(group_id), M, n) + out_ft["res"][group_id] = self._accumulate_sum( + out_ft["res"].get(group_id), M + ) + out_ft_counts["res"][group_id] = out_ft_counts["res"].get(group_id, 0) + 1 def _process_polymer( self, @@ -324,8 +341,9 @@ def _process_polymer( is_highest: bool, out_force: dict[str, dict[Any, Matrix]], out_torque: dict[str, dict[Any, Matrix]], + out_counts: dict[str, dict[Any, int]], out_ft: dict[str, dict[Any, Matrix]] | None, - molcount: dict[int, int], + out_ft_counts: dict[str, dict[Any, int]] | None, combined: bool, ) -> None: """Compute polymer-level force/torque (and optional FT) moments for one @@ -333,10 +351,10 @@ def _process_polymer( Polymer level uses a single bead. Translation/rotation axes, center, and principal moments of inertia are computed, then used to build the - generalized force and torque vectors. Outputs are incrementally averaged - across molecules in the same group for this frame. If combined FT matrices - are enabled and this is the highest level, a force-torque block matrix is - also constructed and averaged. + generalized force and torque vectors. Outputs are accumulated + as sums and counts across molecules in the same group for this frame. + If combined FT matrices are enabled and this is the highest level, + a force-torque block matrix is also constructed and averaged. Args: u: MDAnalysis Universe (or compatible) providing atom access. @@ -350,12 +368,13 @@ def _process_polymer( is_highest: Whether polymer level is the highest level for the molecule. out_force: Output accumulator for polymer force second moments. out_torque: Output accumulator for polymer torque second moments. + out_counts: Output accumulator for polymer molecule counts. out_ft: Optional output accumulator for polymer combined FT matrices. - molcount: Per-group molecule counter for within-frame averaging. + out_ft_counts: Optional output accumulator for polymer FT counts. combined: Whether combined force-torque matrices are enabled. Returns: - None. Mutates output dictionaries and molcount in-place. + None. Mutates output dictionaries and count accumulators in-place. """ bead_key = (mol_id, "polymer") bead_idx_list = beads.get(bead_key, []) @@ -394,20 +413,20 @@ def _process_polymer( F, T = self._ft.compute_frame_covariance(force_vecs, torque_vecs) - n = molcount.get(group_id, 0) + 1 - out_force["poly"][group_id] = self._inc_mean( - out_force["poly"].get(group_id), F, n + out_force["poly"][group_id] = self._accumulate_sum( + out_force["poly"].get(group_id), F ) - out_torque["poly"][group_id] = self._inc_mean( - out_torque["poly"].get(group_id), T, n + out_torque["poly"][group_id] = self._accumulate_sum( + out_torque["poly"].get(group_id), T ) - molcount[group_id] = n + out_counts["poly"][group_id] = out_counts["poly"].get(group_id, 0) + 1 - if combined and is_highest and out_ft is not None: + if combined and is_highest and out_ft is not None and out_ft_counts is not None: M = self._build_ft_block(force_vecs, torque_vecs) - out_ft["poly"][group_id] = self._inc_mean( - out_ft["poly"].get(group_id), M, n + out_ft["poly"][group_id] = self._accumulate_sum( + out_ft["poly"].get(group_id), M ) + out_ft_counts["poly"][group_id] = out_ft_counts["poly"].get(group_id, 0) + 1 def _build_ua_vectors( self, @@ -641,20 +660,19 @@ def _try_get_box(u: Any) -> np.ndarray | None: return None @staticmethod - def _inc_mean(old: np.ndarray | None, new: np.ndarray, n: int) -> np.ndarray: - """Compute an incremental mean (streaming average). + def _accumulate_sum(old: np.ndarray | None, new: np.ndarray) -> np.ndarray: + """Accumulate a deterministic sum of matrix contributions. Args: - old: Previous running mean value, or None for the first sample. - new: New sample to incorporate. - n: 1-based sample count after adding the new sample. + old: Previous running sum value, or None for the first sample. + new: New sample to add into the sum. Returns: - Updated running mean. + Updated running sum. """ if old is None: return new.copy() - return old + (new - old) / float(n) + return old + new @staticmethod def _build_ft_block( diff --git a/CodeEntropy/results/reporter.py b/CodeEntropy/results/reporter.py index 3261c25..fcd2ef2 100644 --- a/CodeEntropy/results/reporter.py +++ b/CodeEntropy/results/reporter.py @@ -408,10 +408,13 @@ def _build_grouped_payload( key = f"{level}:{typ}" groups[gid]["components"][key] = val - for g in groups.values(): + for gid in sorted(groups.keys()): + g = groups[gid] if g["total"] is None: comps = sorted(g["components"].values()) - g["total"] = float(sum(comps)) if comps else 0.0 + g["total"] = ( + float(sum(float(x) for x in sorted(comps))) if comps else 0.0 + ) payload: dict[str, Any] = { "args": self._serialize_args(args), diff --git a/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py index 1b51006..a79a2ff 100644 --- a/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py +++ b/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py @@ -3,7 +3,6 @@ import numpy as np import pytest -from CodeEntropy.levels.nodes import covariance as covmod from CodeEntropy.levels.nodes.covariance import FrameCovarianceNode @@ -42,21 +41,21 @@ def test_try_get_box_returns_none_on_failure(): assert node._try_get_box(u) is None -def test_inc_mean_first_sample_copies(): +def test_accumulate_sum_first_sample_copies(): node = FrameCovarianceNode() new = np.eye(2) - out = node._inc_mean(None, new, n=1) + out = node._accumulate_sum(None, new) np.testing.assert_allclose(out, new) new[0, 0] = 999.0 assert out[0, 0] != 999.0 -def test_inc_mean_updates_streaming_average(): +def test_accumulate_sum_adds_arrays(): node = FrameCovarianceNode() old = np.array([[2.0, 2.0], [2.0, 2.0]]) new = np.array([[4.0, 0.0], [0.0, 4.0]]) - out = node._inc_mean(old, new, n=2) - np.testing.assert_allclose(out, np.array([[3.0, 1.0], [1.0, 3.0]])) + out = node._accumulate_sum(old, new) + np.testing.assert_allclose(out, np.array([[6.0, 2.0], [2.0, 6.0]])) def test_build_ft_block_rejects_mismatched_lengths(): @@ -108,6 +107,8 @@ def test_process_residue_skips_when_no_beads_key_present(): assert out["force"]["res"] == {} assert out["torque"]["res"] == {} assert "forcetorque" not in out + assert out["force_counts"]["res"] == {} + assert out["torque_counts"]["res"] == {} def test_process_residue_combined_only_when_highest_level(): @@ -163,6 +164,9 @@ def test_process_residue_combined_only_when_highest_level(): assert 7 in out["force"]["res"] assert 7 in out["torque"]["res"] assert 7 in out["forcetorque"]["res"] + assert out["force_counts"]["res"][7] == 1 + assert out["torque_counts"]["res"][7] == 1 + assert out["forcetorque_counts"]["res"][7] == 1 def test_process_residue_combined_not_added_if_not_highest_level(): @@ -214,6 +218,8 @@ def test_process_residue_combined_not_added_if_not_highest_level(): assert "forcetorque" in out assert out["forcetorque"]["res"] == {} + assert out["force_counts"]["res"][7] == 1 + assert out["torque_counts"]["res"][7] == 1 def test_process_united_atom_returns_when_no_beads_for_level(): @@ -228,7 +234,7 @@ def test_process_united_atom_returns_when_no_beads_for_level(): out_force = {"ua": {}, "res": {}, "poly": {}} out_torque = {"ua": {}, "res": {}, "poly": {}} - molcount = {} + out_counts = {"ua": {}, "res": {}, "poly": {}} node._process_united_atom( u=MagicMock(), @@ -243,12 +249,12 @@ def test_process_united_atom_returns_when_no_beads_for_level(): is_highest=True, out_force=out_force, out_torque=out_torque, - molcount=molcount, + out_counts=out_counts, ) assert out_force["ua"] == {} assert out_torque["ua"] == {} - assert molcount == {} + assert out_counts["ua"] == {} axes_manager.get_UA_axes.assert_not_called() axes_manager.get_vanilla_axes.assert_not_called() @@ -319,7 +325,7 @@ def test_get_polymer_axes_returns_arrays(monkeypatch): assert np.allclose(moi, np.array([1.0, 1.0, 1.0])) -def test_process_united_atom_updates_outputs_and_molcount(): +def test_process_united_atom_updates_outputs_and_counts(): node = FrameCovarianceNode() node._build_ua_vectors = MagicMock( @@ -345,7 +351,7 @@ def test_process_united_atom_updates_outputs_and_molcount(): beads = {(0, "united_atom", 0): [123]} out_force = {"ua": {}, "res": {}, "poly": {}} out_torque = {"ua": {}, "res": {}, "poly": {}} - molcount = {} + out_counts = {"ua": {}, "res": {}, "poly": {}} node._process_united_atom( u=u, @@ -360,13 +366,13 @@ def test_process_united_atom_updates_outputs_and_molcount(): is_highest=True, out_force=out_force, out_torque=out_torque, - molcount=molcount, + out_counts=out_counts, ) key = (7, 0) assert np.allclose(out_force["ua"][key], F) assert np.allclose(out_torque["ua"][key], T) - assert molcount[key] == 1 + assert out_counts["ua"][key] == 1 def test_process_residue_returns_early_when_no_beads(): @@ -374,6 +380,7 @@ def test_process_residue_returns_early_when_no_beads(): out_force = {"ua": {}, "res": {}, "poly": {}} out_torque = {"ua": {}, "res": {}, "poly": {}} + out_counts = {"ua": {}, "res": {}, "poly": {}} node._process_residue( u=MagicMock(), @@ -388,102 +395,39 @@ def test_process_residue_returns_early_when_no_beads(): is_highest=True, out_force=out_force, out_torque=out_torque, + out_counts=out_counts, out_ft=None, - molcount={}, + out_ft_counts=None, combined=False, ) assert out_force["res"] == {} assert out_torque["res"] == {} - - -def test_build_ua_vectors_customised_axes_true_calls_get_UA_axes(): - node = FrameCovarianceNode() - - bead = _BeadGroup(1) - residue_atoms = MagicMock() - - axes_manager = MagicMock() - axes_manager.get_UA_axes.return_value = ( - np.eye(3), - np.eye(3), - np.array([0.0, 0.0, 0.0]), - np.array([1.0, 1.0, 1.0]), - ) - - node._ft.get_weighted_forces = MagicMock(return_value=np.array([1.0, 2.0, 3.0])) - node._ft.get_weighted_torques = MagicMock(return_value=np.array([4.0, 5.0, 6.0])) - - force_vecs, torque_vecs = node._build_ua_vectors( - bead_groups=[bead], - residue_atoms=residue_atoms, - axes_manager=axes_manager, - box=np.array([10.0, 10.0, 10.0]), - force_partitioning=1.0, - customised_axes=True, - is_highest=True, - ) - - axes_manager.get_UA_axes.assert_called_once() - assert len(force_vecs) == 1 and len(torque_vecs) == 1 - - -def test_build_ua_vectors_vanilla_path_uses_principal_axes_and_vanilla_axes( - monkeypatch, -): - node = FrameCovarianceNode() - - residue_atoms = MagicMock() - residue_atoms.principal_axes.return_value = np.eye(3) - - bead = _BeadGroup(1) - - axes_manager = MagicMock() - axes_manager.get_vanilla_axes.return_value = ( - np.eye(3) * 2, - np.array([9.0, 8.0, 7.0]), - ) - - monkeypatch.setattr(covmod, "make_whole", lambda *_: None) - - node._ft.get_weighted_forces = MagicMock(return_value=np.array([1.0, 0.0, 0.0])) - node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0])) - - force_vecs, torque_vecs = node._build_ua_vectors( - bead_groups=[bead], - residue_atoms=residue_atoms, - axes_manager=axes_manager, - box=np.array([10.0, 10.0, 10.0]), - force_partitioning=1.0, - customised_axes=False, - is_highest=True, - ) - - axes_manager.get_vanilla_axes.assert_called_once() - assert len(force_vecs) == 1 and len(torque_vecs) == 1 + assert out_counts["res"] == {} def test_process_united_atom_skips_when_any_bead_group_is_empty(): node = FrameCovarianceNode() + u = MagicMock() + u.atoms = MagicMock() + u.atoms.__getitem__.side_effect = lambda idx: _EmptyGroup() + res = MagicMock() res.atoms = MagicMock() mol = MagicMock() mol.residues = [res] - u = MagicMock() - u.atoms = MagicMock() - u.atoms.__getitem__.side_effect = lambda idx: _EmptyGroup() - out_force = {"ua": {}, "res": {}, "poly": {}} out_torque = {"ua": {}, "res": {}, "poly": {}} + out_counts = {"ua": {}, "res": {}, "poly": {}} node._process_united_atom( u=u, mol=mol, mol_id=0, group_id=0, - beads={(0, "united_atom", 0): [123]}, + beads={(0, "united_atom", 0): [1]}, axes_manager=MagicMock(), box=np.array([10.0, 10.0, 10.0]), force_partitioning=1.0, @@ -491,11 +435,12 @@ def test_process_united_atom_skips_when_any_bead_group_is_empty(): is_highest=True, out_force=out_force, out_torque=out_torque, - molcount={}, + out_counts=out_counts, ) assert out_force["ua"] == {} assert out_torque["ua"] == {} + assert out_counts["ua"] == {} def test_process_residue_returns_early_when_any_bead_group_is_empty(): @@ -507,13 +452,14 @@ def test_process_residue_returns_early_when_any_bead_group_is_empty(): out_force = {"ua": {}, "res": {}, "poly": {}} out_torque = {"ua": {}, "res": {}, "poly": {}} + out_counts = {"ua": {}, "res": {}, "poly": {}} node._process_residue( u=u, mol=MagicMock(), mol_id=0, group_id=0, - beads={(0, "residue"): [np.array([1, 2, 3])]}, + beads={(0, "residue"): [1]}, axes_manager=MagicMock(), box=np.array([10.0, 10.0, 10.0]), customised_axes=False, @@ -521,13 +467,15 @@ def test_process_residue_returns_early_when_any_bead_group_is_empty(): is_highest=True, out_force=out_force, out_torque=out_torque, - out_ft=None, - molcount={}, + out_counts=out_counts, + out_ft={}, + out_ft_counts={}, combined=False, ) assert out_force["res"] == {} assert out_torque["res"] == {} + assert out_counts["res"] == {} def test_process_polymer_skips_when_any_bead_group_is_empty(): @@ -539,96 +487,187 @@ def test_process_polymer_skips_when_any_bead_group_is_empty(): out_force = {"ua": {}, "res": {}, "poly": {}} out_torque = {"ua": {}, "res": {}, "poly": {}} - out_ft = {"ua": {}, "res": {}, "poly": {}} + out_counts = {"ua": {}, "res": {}, "poly": {}} node._process_polymer( u=u, mol=MagicMock(), mol_id=0, - group_id=7, - beads={(0, "polymer"): [np.array([1, 2, 3])]}, + group_id=0, + beads={(0, "polymer"): [1]}, axes_manager=MagicMock(), box=np.array([10.0, 10.0, 10.0]), force_partitioning=1.0, is_highest=True, out_force=out_force, out_torque=out_torque, - out_ft=out_ft, - molcount={}, - combined=True, + out_counts=out_counts, + out_ft=None, + out_ft_counts=None, + combined=False, ) assert out_force["poly"] == {} assert out_torque["poly"] == {} - assert out_ft["poly"] == {} + assert out_counts["poly"] == {} def test_process_polymer_happy_path_updates_force_torque_and_optional_ft(): node = FrameCovarianceNode() + node._get_polymer_axes = MagicMock( + return_value=(np.eye(3), np.eye(3), np.zeros(3), np.ones(3)) + ) + node._build_ft_block = MagicMock(return_value=np.eye(6)) + node._ft.get_weighted_forces = MagicMock(return_value=np.array([1.0, 0.0, 0.0])) + node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0])) + node._ft.compute_frame_covariance = MagicMock( + return_value=(np.eye(3), 2.0 * np.eye(3)) + ) + u = MagicMock() u.atoms = MagicMock() + bead = _BeadGroup(1) + u.atoms.__getitem__.side_effect = lambda idx: bead - bead_obj = _BeadGroup(1) - u.atoms.__getitem__.side_effect = lambda idx: bead_obj + out_force = {"ua": {}, "res": {}, "poly": {}} + out_torque = {"ua": {}, "res": {}, "poly": {}} + out_counts = {"ua": {}, "res": {}, "poly": {}} + out_ft = {"ua": {}, "res": {}, "poly": {}} + out_ft_counts = {"ua": {}, "res": {}, "poly": {}} - mol = MagicMock() - mol.atoms = MagicMock() + node._process_polymer( + u=u, + mol=MagicMock(), + mol_id=0, + group_id=7, + beads={(0, "polymer"): [1]}, + axes_manager=MagicMock(), + box=np.array([10.0, 10.0, 10.0]), + force_partitioning=1.0, + is_highest=True, + out_force=out_force, + out_torque=out_torque, + out_counts=out_counts, + out_ft=out_ft, + out_ft_counts=out_ft_counts, + combined=True, + ) + + assert 7 in out_force["poly"] + assert 7 in out_torque["poly"] + assert 7 in out_ft["poly"] + assert out_counts["poly"][7] == 1 + assert out_ft_counts["poly"][7] == 1 + + +def test_build_ua_vectors_customised_axes_uses_get_UA_axes(): + node = FrameCovarianceNode() + + bead0 = MagicMock() + bead1 = MagicMock() + bead_groups = [bead0, bead1] + residue_atoms = MagicMock() axes_manager = MagicMock() + axes_manager.get_UA_axes.side_effect = [ + ( + np.eye(3), + np.eye(3) * 2, + np.array([1.0, 2.0, 3.0]), + np.array([4.0, 5.0, 6.0]), + ), + ( + np.eye(3) * 3, + np.eye(3) * 4, + np.array([7.0, 8.0, 9.0]), + np.array([10.0, 11.0, 12.0]), + ), + ] - f_vec = np.array([1.0, 0.0, 0.0], dtype=float) - t_vec = np.array([0.0, 1.0, 0.0], dtype=float) + node._ft.get_weighted_forces = MagicMock( + side_effect=[np.array([1.0, 0.0, 0.0]), np.array([2.0, 0.0, 0.0])] + ) + node._ft.get_weighted_torques = MagicMock( + side_effect=[np.array([0.0, 1.0, 0.0]), np.array([0.0, 2.0, 0.0])] + ) - F = np.eye(3) - T = 2.0 * np.eye(3) - FT = np.eye(6) + force_vecs, torque_vecs = node._build_ua_vectors( + bead_groups=bead_groups, + residue_atoms=residue_atoms, + axes_manager=axes_manager, + box=np.array([10.0, 10.0, 10.0]), + force_partitioning=1.5, + customised_axes=True, + is_highest=True, + ) - out_force = {"ua": {}, "res": {}, "poly": {}} - out_torque = {"ua": {}, "res": {}, "poly": {}} - out_ft = {"ua": {}, "res": {}, "poly": {}} - molcount = {} + assert len(force_vecs) == 2 + assert len(torque_vecs) == 2 - with ( - patch.object( - node, - "_get_polymer_axes", - return_value=(np.eye(3), np.eye(3), np.zeros(3), np.ones(3)), - ) as axes_spy, - patch.object(node._ft, "get_weighted_forces", return_value=f_vec) as f_spy, - patch.object(node._ft, "get_weighted_torques", return_value=t_vec) as t_spy, - patch.object( - node._ft, "compute_frame_covariance", return_value=(F, T) - ) as cov_spy, - patch.object(node, "_build_ft_block", return_value=FT) as ft_spy, - ): - node._process_polymer( - u=u, - mol=mol, - mol_id=0, - group_id=7, - beads={(0, "polymer"): [np.array([1, 2, 3])]}, + np.testing.assert_allclose(force_vecs[0], np.array([1.0, 0.0, 0.0])) + np.testing.assert_allclose(force_vecs[1], np.array([2.0, 0.0, 0.0])) + np.testing.assert_allclose(torque_vecs[0], np.array([0.0, 1.0, 0.0])) + np.testing.assert_allclose(torque_vecs[1], np.array([0.0, 2.0, 0.0])) + + assert axes_manager.get_UA_axes.call_count == 2 + axes_manager.get_UA_axes.assert_any_call(residue_atoms, 0) + axes_manager.get_UA_axes.assert_any_call(residue_atoms, 1) + + assert node._ft.get_weighted_forces.call_count == 2 + assert node._ft.get_weighted_torques.call_count == 2 + + +def test_build_ua_vectors_vanilla_axes_uses_make_whole_and_vanilla_axes(): + node = FrameCovarianceNode() + + bead0 = MagicMock() + bead1 = MagicMock() + bead0.center_of_mass.return_value = np.array([1.0, 1.0, 1.0]) + bead1.center_of_mass.return_value = np.array([2.0, 2.0, 2.0]) + + bead_groups = [bead0, bead1] + residue_atoms = MagicMock() + residue_atoms.principal_axes.return_value = np.eye(3) + + axes_manager = MagicMock() + axes_manager.get_vanilla_axes.side_effect = [ + (np.eye(3) * 2, np.array([3.0, 4.0, 5.0])), + (np.eye(3) * 3, np.array([6.0, 7.0, 8.0])), + ] + + node._ft.get_weighted_forces = MagicMock( + side_effect=[np.array([1.0, 0.0, 0.0]), np.array([2.0, 0.0, 0.0])] + ) + node._ft.get_weighted_torques = MagicMock( + side_effect=[np.array([0.0, 1.0, 0.0]), np.array([0.0, 2.0, 0.0])] + ) + + with patch("CodeEntropy.levels.nodes.covariance.make_whole") as make_whole: + force_vecs, torque_vecs = node._build_ua_vectors( + bead_groups=bead_groups, + residue_atoms=residue_atoms, axes_manager=axes_manager, - box=np.array([10.0, 10.0, 10.0]), - force_partitioning=0.5, - is_highest=True, - out_force=out_force, - out_torque=out_torque, - out_ft=out_ft, - molcount=molcount, - combined=True, + box=np.array([20.0, 20.0, 20.0]), + force_partitioning=2.0, + customised_axes=False, + is_highest=False, ) - assert u.atoms.__getitem__.call_count == 1 - axes_spy.assert_called_once_with(mol=mol, bead=bead_obj, axes_manager=axes_manager) + assert len(force_vecs) == 2 + assert len(torque_vecs) == 2 - f_spy.assert_called_once() - t_spy.assert_called_once() - cov_spy.assert_called_once() + assert make_whole.call_count == 4 + make_whole.assert_any_call(residue_atoms) + make_whole.assert_any_call(bead0) + make_whole.assert_any_call(bead1) - np.testing.assert_allclose(out_force["poly"][7], F) - np.testing.assert_allclose(out_torque["poly"][7], T) - assert molcount[7] == 1 + assert residue_atoms.principal_axes.call_count == 2 + assert axes_manager.get_vanilla_axes.call_count == 2 + bead0.center_of_mass.assert_called_once_with(unwrap=True) + bead1.center_of_mass.assert_called_once_with(unwrap=True) - ft_spy.assert_called_once() - np.testing.assert_allclose(out_ft["poly"][7], FT) + np.testing.assert_allclose(force_vecs[0], np.array([1.0, 0.0, 0.0])) + np.testing.assert_allclose(force_vecs[1], np.array([2.0, 0.0, 0.0])) + np.testing.assert_allclose(torque_vecs[0], np.array([0.0, 1.0, 0.0])) + np.testing.assert_allclose(torque_vecs[1], np.array([0.0, 2.0, 0.0])) diff --git a/tests/unit/CodeEntropy/levels/nodes/test_init_covariance_accumulators_node.py b/tests/unit/CodeEntropy/levels/nodes/test_init_covariance_accumulators_node.py index fd5d31e..a8f1487 100644 --- a/tests/unit/CodeEntropy/levels/nodes/test_init_covariance_accumulators_node.py +++ b/tests/unit/CodeEntropy/levels/nodes/test_init_covariance_accumulators_node.py @@ -10,14 +10,66 @@ def test_init_covariance_accumulators_allocates_and_sets_aliases(): out = node.run(shared) - assert out["group_id_to_index"] == {9: 0, 2: 1} - assert out["index_to_group_id"] == [9, 2] + gid2i = out["group_id_to_index"] + i2gid = out["index_to_group_id"] - assert shared["force_covariances"]["res"] == [None, None] - assert shared["torque_covariances"]["poly"] == [None, None] + assert set(gid2i.keys()) == {9, 2} + assert set(i2gid) == {9, 2} - assert np.all(shared["frame_counts"]["res"] == np.array([0, 0])) - assert np.all(shared["forcetorque_counts"]["poly"] == np.array([0, 0])) + for gid, idx in gid2i.items(): + assert i2gid[idx] == gid - assert shared["force_torque_stats"] is shared["forcetorque_covariances"] - assert shared["force_torque_counts"] is shared["forcetorque_counts"] + assert sorted(gid2i.values()) == [0, 1] + + assert "force_sums" in out + assert "torque_sums" in out + assert "force_counts" in out + assert "torque_counts" in out + assert "forcetorque_sums" in out + assert "forcetorque_counts" in out + + assert "force_covariances" in out + assert "torque_covariances" in out + assert "frame_counts" in out + assert "forcetorque_covariances" in out + assert "force_torque_stats" in out + assert "force_torque_counts" in out + + assert out["force_covariances"] is out["force_sums"] + assert out["torque_covariances"] is out["torque_sums"] + assert out["forcetorque_covariances"] is out["forcetorque_sums"] + assert out["frame_counts"] is out["force_counts"] + + assert out["force_torque_stats"] is out["forcetorque_sums"] + assert out["force_torque_counts"] is out["forcetorque_counts"] + + +def test_init_covariance_accumulators_is_fully_deterministic(): + node = InitCovarianceAccumulatorsNode() + + shared1 = {"groups": {9: [1, 2], 2: [3]}} + shared2 = {"groups": {2: [3], 9: [1, 2]}} + + out1 = node.run(shared1.copy()) + out2 = node.run(shared2.copy()) + + assert out1["group_id_to_index"] == out2["group_id_to_index"] + assert out1["index_to_group_id"] == out2["index_to_group_id"] + + +def test_init_covariance_accumulators_aliases_are_intentional(): + node = InitCovarianceAccumulatorsNode() + + shared = {"groups": {1: [1]}} + out = node.run(shared) + + assert out["force_covariances"] is out["force_sums"] + assert out["torque_covariances"] is out["torque_sums"] + assert out["forcetorque_covariances"] is out["forcetorque_sums"] + assert out["frame_counts"] is out["force_counts"] + assert out["force_torque_stats"] is out["forcetorque_sums"] + assert out["force_torque_counts"] is out["forcetorque_counts"] + + assert np.array_equal(out["force_counts"]["res"], np.array([0])) + assert np.array_equal(out["torque_counts"]["res"], np.array([0])) + assert np.array_equal(out["forcetorque_counts"]["res"], np.array([0])) diff --git a/tests/unit/CodeEntropy/levels/test_level_dag_orchestration.py b/tests/unit/CodeEntropy/levels/test_level_dag_orchestration.py index 17a5a93..f3e28d8 100644 --- a/tests/unit/CodeEntropy/levels/test_level_dag_orchestration.py +++ b/tests/unit/CodeEntropy/levels/test_level_dag_orchestration.py @@ -8,16 +8,27 @@ def _shared(): return { "levels": [["united_atom"]], - "frame_counts": {}, + "group_id_to_index": {0: 0}, + "force_sums": {"ua": {}, "res": [None], "poly": [None]}, + "torque_sums": {"ua": {}, "res": [None], "poly": [None]}, + "force_counts": { + "ua": {}, + "res": np.zeros(1, dtype=int), + "poly": np.zeros(1, dtype=int), + }, + "torque_counts": { + "ua": {}, + "res": np.zeros(1, dtype=int), + "poly": np.zeros(1, dtype=int), + }, + "forcetorque_sums": {"res": [None], "poly": [None]}, + "forcetorque_counts": { + "res": np.zeros(1, dtype=int), + "poly": np.zeros(1, dtype=int), + }, "force_covariances": {}, "torque_covariances": {}, - "force_counts": {}, - "torque_counts": {}, - "reduced_force_covariances": {}, - "reduced_torque_covariances": {}, - "reduced_force_counts": {}, - "reduced_torque_counts": {}, - "group_id_to_index": {0: 0}, + "forcetorque_covariances": {}, } @@ -29,16 +40,35 @@ def test_execute_sets_default_axes_manager_once(): "start": 0, "end": 0, "step": 1, + "force_sums": {"ua": {}, "res": [None], "poly": [None]}, + "torque_sums": {"ua": {}, "res": [None], "poly": [None]}, + "force_counts": { + "ua": {}, + "res": np.zeros(1, dtype=int), + "poly": np.zeros(1, dtype=int), + }, + "torque_counts": { + "ua": {}, + "res": np.zeros(1, dtype=int), + "poly": np.zeros(1, dtype=int), + }, + "forcetorque_sums": {"res": [None], "poly": [None]}, + "forcetorque_counts": { + "res": np.zeros(1, dtype=int), + "poly": np.zeros(1, dtype=int), + }, } dag._run_static_stage = MagicMock() dag._run_frame_stage = MagicMock() + dag._finalize_means = MagicMock() dag.execute(shared) assert "axes_manager" in shared dag._run_static_stage.assert_called_once() dag._run_frame_stage.assert_called_once() + dag._finalize_means.assert_called_once_with(shared) def test_run_static_stage_calls_nodes_in_topological_sort_order(): @@ -71,6 +101,8 @@ def test_run_frame_stage_iterates_selected_frames_and_reduces_each(): { "force": {"ua": {}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, } ] * 2 dag._reduce_one_frame = MagicMock() @@ -83,15 +115,10 @@ def test_run_frame_stage_iterates_selected_frames_and_reduces_each(): dag._frame_dag.execute_frame.assert_any_call(shared, 11) -def test_incremental_mean_handles_non_copyable_values(): - out = LevelDAG._incremental_mean(old=None, new=3.0, n=1) - assert out == 3.0 - - def test_reduce_forcetorque_no_key_is_noop(): dag = LevelDAG() shared = { - "forcetorque_covariances": {"res": [None], "poly": [None]}, + "forcetorque_sums": {"res": [None], "poly": [None]}, "forcetorque_counts": { "res": np.zeros(1, dtype=int), "poly": np.zeros(1, dtype=int), @@ -100,7 +127,7 @@ def test_reduce_forcetorque_no_key_is_noop(): } dag._reduce_forcetorque(shared, frame_out={}) assert shared["forcetorque_counts"]["res"][0] == 0 - assert shared["forcetorque_covariances"]["res"][0] is None + assert shared["forcetorque_sums"]["res"][0] is None def test_build_registers_static_nodes_and_builds_frame_dag(): @@ -136,9 +163,10 @@ def test_reduce_force_and_torque_hits_zero_count_branches(): dag = LevelDAG() shared = { - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": {"ua": {}, "res": [0], "poly": [0]}, + "force_sums": {"ua": {}, "res": [None], "poly": [None]}, + "torque_sums": {"ua": {}, "res": [None], "poly": [None]}, + "force_counts": {"ua": {}, "res": np.array([0]), "poly": np.array([0])}, + "torque_counts": {"ua": {}, "res": np.array([0]), "poly": np.array([0])}, "group_id_to_index": {7: 0}, } @@ -153,16 +181,20 @@ def test_reduce_force_and_torque_hits_zero_count_branches(): "res": {7: np.eye(2)}, "poly": {7: np.eye(3)}, }, + "force_counts": {"ua": {(7, 0): 1}, "res": {7: 1}, "poly": {7: 1}}, + "torque_counts": {"ua": {(7, 0): 1}, "res": {7: 1}, "poly": {7: 1}}, } dag._reduce_force_and_torque(shared, frame_out) - assert shared["frame_counts"]["ua"][(7, 0)] == 1 - assert (7, 0) in shared["force_covariances"]["ua"] - assert (7, 0) in shared["torque_covariances"]["ua"] + assert shared["force_counts"]["ua"][(7, 0)] == 1 + assert (7, 0) in shared["force_sums"]["ua"] + assert (7, 0) in shared["torque_sums"]["ua"] - assert shared["frame_counts"]["res"][0] == 1 - assert shared["frame_counts"]["poly"][0] == 1 + assert shared["force_counts"]["res"][0] == 1 + assert shared["force_counts"]["poly"][0] == 1 + assert shared["torque_counts"]["res"][0] == 1 + assert shared["torque_counts"]["poly"][0] == 1 def test_reduce_force_and_torque_handles_empty_frame_gracefully(): @@ -170,22 +202,27 @@ def test_reduce_force_and_torque_handles_empty_frame_gracefully(): shared = { "group_id_to_index": {0: 0}, - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": {"ua": {}, "res": [0], "poly": [0]}, + "force_sums": {"ua": {}, "res": [None], "poly": [None]}, + "torque_sums": {"ua": {}, "res": [None], "poly": [None]}, + "force_counts": {"ua": {}, "res": np.array([0]), "poly": np.array([0])}, + "torque_counts": {"ua": {}, "res": np.array([0]), "poly": np.array([0])}, } frame_out = { "force": {"ua": {}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, } dag._reduce_force_and_torque(shared_data=shared, frame_out=frame_out) - assert shared["force_covariances"]["ua"] == {} - assert shared["torque_covariances"]["ua"] == {} - assert shared["frame_counts"]["res"][0] == 0 - assert shared["frame_counts"]["poly"][0] == 0 + assert shared["force_sums"]["ua"] == {} + assert shared["torque_sums"]["ua"] == {} + assert shared["force_counts"]["res"][0] == 0 + assert shared["force_counts"]["poly"][0] == 0 + assert shared["torque_counts"]["res"][0] == 0 + assert shared["torque_counts"]["poly"][0] == 0 def test_reduce_force_and_torque_increments_res_and_poly_counts_from_zero(): @@ -193,9 +230,12 @@ def test_reduce_force_and_torque_increments_res_and_poly_counts_from_zero(): shared = { "group_id_to_index": {7: 0}, - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": {"ua": {}, "res": [0], "poly": [0]}, + "force_sums": {"ua": {}, "res": [None], "poly": [None]}, + "torque_sums": {"ua": {}, "res": [None], "poly": [None]}, + "force_counts": {"ua": {}, "res": np.array([0]), "poly": np.array([0])}, + "torque_counts": {"ua": {}, "res": np.array([0]), "poly": np.array([0])}, + "forcetorque_sums": {"res": [None], "poly": [None]}, + "forcetorque_counts": {"res": np.array([0]), "poly": np.array([0])}, } F = np.eye(3) @@ -204,12 +244,24 @@ def test_reduce_force_and_torque_increments_res_and_poly_counts_from_zero(): frame_out = { "force": {"ua": {}, "res": {7: F}, "poly": {7: F}}, "torque": {"ua": {}, "res": {7: T}, "poly": {7: T}}, + "force_counts": {"ua": {}, "res": {7: 1}, "poly": {7: 1}}, + "torque_counts": {"ua": {}, "res": {7: 1}, "poly": {7: 1}}, } dag._reduce_force_and_torque(shared_data=shared, frame_out=frame_out) - assert shared["frame_counts"]["res"][0] == 1 - assert shared["frame_counts"]["poly"][0] == 1 + assert shared["force_counts"]["res"][0] == 1 + assert shared["force_counts"]["poly"][0] == 1 + assert shared["torque_counts"]["res"][0] == 1 + assert shared["torque_counts"]["poly"][0] == 1 + assert np.allclose(shared["torque_sums"]["res"][0], T) + assert np.allclose(shared["torque_sums"]["poly"][0], T) + + shared["force_covariances"] = {} + shared["torque_covariances"] = {} + shared["forcetorque_covariances"] = {} + dag._finalize_means(shared) + assert np.allclose(shared["torque_covariances"]["res"][0], T) assert np.allclose(shared["torque_covariances"]["poly"][0], T) @@ -218,43 +270,42 @@ def test_reduce_one_frame_skips_missing_force_and_torque_keys(): dag = LevelDAG() shared = _shared() - bead_key = (0, "united_atom", 0) frame_out = { - "beads": {bead_key: [1, 2, 3]}, - "counts": {bead_key: 1}, "force": {"ua": {}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, } dag._reduce_one_frame(shared_data=shared, frame_out=frame_out) - assert shared["force_covariances"] == {} - assert shared["torque_covariances"] == {} + assert shared["force_sums"]["ua"] == {} + assert shared["torque_sums"]["ua"] == {} def test_reduce_force_and_torque_skips_when_counts_are_zero(): dag = LevelDAG() shared = _shared() - k = (0, "united_atom", 0) - shared["force_covariances"][k] = np.eye(3) - shared["torque_covariances"][k] = np.eye(3) - shared["force_counts"][k] = 0 - shared["torque_counts"][k] = 0 - shared["frame_counts"][k] = 0 + k = (0, 0) + shared["force_sums"]["ua"][k] = np.eye(3) + shared["torque_sums"]["ua"][k] = np.eye(3) + shared["force_counts"]["ua"][k] = 0 + shared["torque_counts"]["ua"][k] = 0 frame_out = { - "force": {"ua": {}, "res": {}, "poly": {}}, - "torque": {"ua": {}, "res": {}, "poly": {}}, - "beads": {}, + "force": {"ua": {k: np.eye(3)}, "res": {}, "poly": {}}, + "torque": {"ua": {k: np.eye(3)}, "res": {}, "poly": {}}, + "force_counts": {"ua": {k: 0}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {k: 0}, "res": {}, "poly": {}}, } dag._reduce_force_and_torque(shared_data=shared, frame_out=frame_out) - assert shared["reduced_force_covariances"] == {} - assert shared["reduced_torque_covariances"] == {} - assert shared["reduced_force_counts"] == {} - assert shared["reduced_torque_counts"] == {} + np.testing.assert_array_equal(shared["force_sums"]["ua"][k], np.eye(3)) + np.testing.assert_array_equal(shared["torque_sums"]["ua"][k], np.eye(3)) + assert shared["force_counts"]["ua"][k] == 0 + assert shared["torque_counts"]["ua"][k] == 0 def test_run_static_stage_forwards_progress_when_node_accepts_it(): @@ -284,7 +335,7 @@ def run(self, shared_data): progress = MagicMock() with patch("networkx.topological_sort", return_value=["a"]): - dag._run_static_stage({"X": 1}, progress=progress) # should not raise + dag._run_static_stage({"X": 1}, progress=progress) def test_run_frame_stage_with_progress_creates_task_and_updates_titles(): @@ -301,6 +352,8 @@ def test_run_frame_stage_with_progress_creates_task_and_updates_titles(): dag._frame_dag.execute_frame.return_value = { "force": {"ua": {}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, } dag._reduce_one_frame = MagicMock() @@ -333,6 +386,8 @@ def test_run_frame_stage_with_negative_end_computes_total_frames(): dag._frame_dag.execute_frame.return_value = { "force": {"ua": {}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, } dag._reduce_one_frame = MagicMock() @@ -349,7 +404,6 @@ def test_run_frame_stage_with_negative_end_computes_total_frames(): def test_run_frame_stage_progress_total_frames_falls_back_to_none_on_error(): - dag = LevelDAG() class BadTrajectory: diff --git a/tests/unit/CodeEntropy/levels/test_level_dag_reduce.py b/tests/unit/CodeEntropy/levels/test_level_dag_reduce.py index cb14bc1..4e2a5b1 100644 --- a/tests/unit/CodeEntropy/levels/test_level_dag_reduce.py +++ b/tests/unit/CodeEntropy/levels/test_level_dag_reduce.py @@ -3,23 +3,25 @@ from CodeEntropy.levels.level_dag import LevelDAG -def test_incremental_mean_first_sample_copies(): - x = np.array([1.0, 2.0]) - out = LevelDAG._incremental_mean(None, x, n=1) - assert np.allclose(out, x) - x[0] = 999.0 - assert out[0] != 999.0 +def _shared(): + return { + "force_sums": {"ua": {}, "res": [None], "poly": [None]}, + "torque_sums": {"ua": {}, "res": [None], "poly": [None]}, + "force_counts": {"ua": {}, "res": np.array([0]), "poly": np.array([0])}, + "torque_counts": {"ua": {}, "res": np.array([0]), "poly": np.array([0])}, + "forcetorque_sums": {"res": [None], "poly": [None]}, + "forcetorque_counts": {"res": np.array([0]), "poly": np.array([0])}, + "force_covariances": {}, + "torque_covariances": {}, + "forcetorque_covariances": {}, + "group_id_to_index": {7: 0}, + } def test_reduce_force_and_torque_exercises_count_branches(): dag = LevelDAG() - shared = { - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": {"ua": {}, "res": [0], "poly": [0]}, - "group_id_to_index": {7: 0}, - } + shared = _shared() frame_out = { "force": { @@ -32,20 +34,24 @@ def test_reduce_force_and_torque_exercises_count_branches(): "res": {7: np.array([5.0])}, "poly": {7: np.array([6.0])}, }, + "force_counts": {"ua": {(9, 0): 1}, "res": {7: 1}, "poly": {7: 1}}, + "torque_counts": {"ua": {(9, 0): 1}, "res": {7: 1}, "poly": {7: 1}}, } dag._reduce_force_and_torque(shared, frame_out) - assert (9, 0) in shared["torque_covariances"]["ua"] - assert shared["frame_counts"]["res"][0] == 1 - assert shared["frame_counts"]["poly"][0] == 1 + assert (9, 0) in shared["torque_sums"]["ua"] + assert shared["force_counts"]["res"][0] == 1 + assert shared["force_counts"]["poly"][0] == 1 + assert shared["torque_counts"]["res"][0] == 1 + assert shared["torque_counts"]["poly"][0] == 1 def test_reduce_forcetorque_returns_when_missing_key(): dag = LevelDAG() shared = { - "forcetorque_covariances": {"res": [None], "poly": [None]}, - "forcetorque_counts": {"res": [0], "poly": [0]}, + "forcetorque_sums": {"res": [None], "poly": [None]}, + "forcetorque_counts": {"res": np.array([0]), "poly": np.array([0])}, "group_id_to_index": {7: 0}, } dag._reduce_forcetorque(shared, frame_out={}) @@ -56,8 +62,8 @@ def test_reduce_forcetorque_updates_res_and_poly(): dag = LevelDAG() shared = { - "forcetorque_covariances": {"res": [None], "poly": [None]}, - "forcetorque_counts": {"res": [0], "poly": [0]}, + "forcetorque_sums": {"res": [None], "poly": [None]}, + "forcetorque_counts": {"res": np.array([0]), "poly": np.array([0])}, "group_id_to_index": {7: 0}, } @@ -65,77 +71,69 @@ def test_reduce_forcetorque_updates_res_and_poly(): "forcetorque": { "res": {7: np.array([1.0, 1.0])}, "poly": {7: np.array([2.0, 2.0])}, - } + }, + "forcetorque_counts": {"res": {7: 1}, "poly": {7: 1}}, } dag._reduce_forcetorque(shared, frame_out) assert shared["forcetorque_counts"]["res"][0] == 1 assert shared["forcetorque_counts"]["poly"][0] == 1 - assert shared["forcetorque_covariances"]["res"][0] is not None - assert shared["forcetorque_covariances"]["poly"][0] is not None + assert shared["forcetorque_sums"]["res"][0] is not None + assert shared["forcetorque_sums"]["poly"][0] is not None def test_reduce_force_and_torque_res_torque_increments_when_res_count_is_zero(): dag = LevelDAG() - shared = { - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": {"ua": {}, "res": [0], "poly": [0]}, - "group_id_to_index": {7: 0}, - } + shared = _shared() frame_out = { "force": {"ua": {}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {7: np.eye(3)}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {7: 1}, "poly": {}}, } dag._reduce_force_and_torque(shared, frame_out) - assert shared["frame_counts"]["res"][0] == 1 - assert shared["torque_covariances"]["res"][0] is not None + assert shared["torque_counts"]["res"][0] == 1 + assert shared["torque_sums"]["res"][0] is not None def test_reduce_force_and_torque_poly_torque_increments_when_poly_count_is_zero(): dag = LevelDAG() - shared = { - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": {"ua": {}, "res": [0], "poly": [0]}, - "group_id_to_index": {7: 0}, - } + shared = _shared() frame_out = { "force": {"ua": {}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {}, "poly": {7: np.eye(3)}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {7: 1}}, } dag._reduce_force_and_torque(shared, frame_out) - assert shared["frame_counts"]["poly"][0] == 1 - assert shared["torque_covariances"]["poly"][0] is not None + assert shared["torque_counts"]["poly"][0] == 1 + assert shared["torque_sums"]["poly"][0] is not None def test_reduce_force_and_torque_increments_ua_frame_counts_for_force(): dag = LevelDAG() - shared = { - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": {"ua": {}, "res": [0], "poly": [0]}, - "group_id_to_index": {7: 0}, - } + shared = _shared() k = (9, 0) frame_out = { "force": {"ua": {k: np.eye(3)}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {k: 1}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, } dag._reduce_force_and_torque(shared, frame_out) - assert shared["frame_counts"]["ua"][k] == 1 - assert k in shared["force_covariances"]["ua"] + assert shared["force_counts"]["ua"][k] == 1 + assert k in shared["force_sums"]["ua"] def test_reduce_force_and_torque_increments_ua_counts_from_zero(): @@ -144,44 +142,37 @@ def test_reduce_force_and_torque_increments_ua_counts_from_zero(): key = (9, 0) F = np.eye(3) - shared = { - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": {"ua": {}, "res": [0], "poly": [0]}, - "group_id_to_index": {7: 0}, - } + shared = _shared() frame_out = { "force": {"ua": {key: F}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {key: 1}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, } dag._reduce_force_and_torque(shared, frame_out) - assert shared["frame_counts"]["ua"][key] == 1 - - np.testing.assert_array_equal(shared["force_covariances"]["ua"][key], F) + assert shared["force_counts"]["ua"][key] == 1 + np.testing.assert_array_equal(shared["force_sums"]["ua"][key], F) def test_reduce_force_and_torque_hits_ua_force_count_increment_line(): dag = LevelDAG() key = (9, 0) - shared = { - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": {"ua": {}, "res": [0], "poly": [0]}, - "group_id_to_index": {7: 0}, - } + shared = _shared() frame_out = { "force": {"ua": {key: np.eye(3)}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {key: 1}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, } dag._reduce_force_and_torque(shared, frame_out) - assert shared["frame_counts"]["ua"][key] == 1 + assert shared["force_counts"]["ua"][key] == 1 def test_reduce_force_and_torque_ua_torque_increments_count_when_force_missing_key(): @@ -190,19 +181,194 @@ def test_reduce_force_and_torque_ua_torque_increments_count_when_force_missing_k key = (9, 0) T = np.eye(3) - shared = { - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": {"ua": {}, "res": [0], "poly": [0]}, - "group_id_to_index": {7: 0}, - } + shared = _shared() frame_out = { "force": {"ua": {}, "res": {}, "poly": {}}, "torque": {"ua": {key: T}, "res": {}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {key: 1}, "res": {}, "poly": {}}, + } + + dag._reduce_force_and_torque(shared, frame_out) + + assert shared["torque_counts"]["ua"][key] == 1 + np.testing.assert_array_equal(shared["torque_sums"]["ua"][key], T) + + +def test_reduce_one_frame_calls_both_reducers(): + dag = LevelDAG() + shared = _shared() + frame_out = {"force": {}, "torque": {}} + + called = {"force": False, "ft": False} + + def fake_reduce_force(shared_data, frame_out_arg): + called["force"] = True + + def fake_reduce_ft(shared_data, frame_out_arg): + called["ft"] = True + + dag._reduce_force_and_torque = fake_reduce_force + dag._reduce_forcetorque = fake_reduce_ft + + dag._reduce_one_frame(shared, frame_out) + + assert called["force"] is True + assert called["ft"] is True + + +def test_reduce_force_and_torque_ua_force_continue_when_count_is_zero(): + dag = LevelDAG() + shared = _shared() + key = (7, 0) + + frame_out = { + "force": {"ua": {key: np.eye(2)}, "res": {}, "poly": {}}, + "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {key: 0}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, + } + + dag._reduce_force_and_torque(shared, frame_out) + + assert key not in shared["force_sums"]["ua"] + assert key not in shared["force_counts"]["ua"] + + +def test_reduce_force_and_torque_ua_torque_continue_when_count_is_negative(): + dag = LevelDAG() + shared = _shared() + key = (7, 0) + + frame_out = { + "force": {"ua": {}, "res": {}, "poly": {}}, + "torque": {"ua": {key: np.eye(2)}, "res": {}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {key: -3}, "res": {}, "poly": {}}, + } + + dag._reduce_force_and_torque(shared, frame_out) + + assert key not in shared["torque_sums"]["ua"] + assert key not in shared["torque_counts"]["ua"] + + +def test_reduce_force_and_torque_res_force_continue_when_count_is_zero(): + dag = LevelDAG() + shared = _shared() + + frame_out = { + "force": {"ua": {}, "res": {7: np.eye(3)}, "poly": {}}, + "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {}, "res": {7: 0}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, + } + + dag._reduce_force_and_torque(shared, frame_out) + + assert shared["force_sums"]["res"][0] is None + assert shared["force_counts"]["res"][0] == 0 + + +def test_reduce_force_and_torque_res_torque_continue_when_count_is_zero(): + dag = LevelDAG() + shared = _shared() + + frame_out = { + "force": {"ua": {}, "res": {}, "poly": {}}, + "torque": {"ua": {}, "res": {7: np.eye(3)}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {7: 0}, "poly": {}}, + } + + dag._reduce_force_and_torque(shared, frame_out) + + assert shared["torque_sums"]["res"][0] is None + assert shared["torque_counts"]["res"][0] == 0 + + +def test_reduce_force_and_torque_poly_force_continue_when_count_is_zero(): + dag = LevelDAG() + shared = _shared() + + frame_out = { + "force": {"ua": {}, "res": {}, "poly": {7: np.eye(3)}}, + "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {7: 0}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, + } + + dag._reduce_force_and_torque(shared, frame_out) + + assert shared["force_sums"]["poly"][0] is None + assert shared["force_counts"]["poly"][0] == 0 + + +def test_reduce_force_and_torque_poly_torque_continue_when_count_is_zero(): + dag = LevelDAG() + shared = _shared() + + frame_out = { + "force": {"ua": {}, "res": {}, "poly": {}}, + "torque": {"ua": {}, "res": {}, "poly": {7: np.eye(3)}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {7: 0}}, + } + + dag._reduce_force_and_torque(shared, frame_out) + + assert shared["torque_sums"]["poly"][0] is None + assert shared["torque_counts"]["poly"][0] == 0 + + +def test_reduce_forcetorque_res_continue_when_count_is_zero(): + dag = LevelDAG() + shared = _shared() + + frame_out = { + "forcetorque": {"res": {7: np.eye(4)}, "poly": {}}, + "forcetorque_counts": {"res": {7: 0}, "poly": {}}, + } + + dag._reduce_forcetorque(shared, frame_out) + + assert shared["forcetorque_sums"]["res"][0] is None + assert shared["forcetorque_counts"]["res"][0] == 0 + + +def test_reduce_forcetorque_poly_continue_when_count_is_zero(): + dag = LevelDAG() + shared = _shared() + + frame_out = { + "forcetorque": {"res": {}, "poly": {7: np.eye(4)}}, + "forcetorque_counts": {"res": {}, "poly": {7: 0}}, + } + + dag._reduce_forcetorque(shared, frame_out) + + assert shared["forcetorque_sums"]["poly"][0] is None + assert shared["forcetorque_counts"]["poly"][0] == 0 + + +def test_reduce_force_and_torque_updates_when_count_is_positive(): + dag = LevelDAG() + shared = _shared() + + F = np.eye(3) + T = np.eye(3) * 2 + + frame_out = { + "force": {"ua": {}, "res": {7: F}, "poly": {}}, + "torque": {"ua": {}, "res": {7: T}, "poly": {}}, + "force_counts": {"ua": {}, "res": {7: 1}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {7: 1}, "poly": {}}, } dag._reduce_force_and_torque(shared, frame_out) - assert shared["frame_counts"]["ua"][key] == 1 - np.testing.assert_array_equal(shared["torque_covariances"]["ua"][key], T) + assert shared["force_counts"]["res"][0] == 1 + assert shared["torque_counts"]["res"][0] == 1 + np.testing.assert_allclose(shared["force_sums"]["res"][0], F) + np.testing.assert_allclose(shared["torque_sums"]["res"][0], T) diff --git a/tests/unit/CodeEntropy/levels/test_level_dag_reduction.py b/tests/unit/CodeEntropy/levels/test_level_dag_reduction.py index 5c8c471..fb8d57e 100644 --- a/tests/unit/CodeEntropy/levels/test_level_dag_reduction.py +++ b/tests/unit/CodeEntropy/levels/test_level_dag_reduction.py @@ -5,32 +5,30 @@ from CodeEntropy.levels.level_dag import LevelDAG -def test_incremental_mean_none_returns_copy_for_numpy(): - arr = np.array([1.0, 2.0]) - out = LevelDAG._incremental_mean(None, arr, n=1) - assert np.all(out == arr) - arr[0] = 999.0 - assert out[0] != 999.0 - - -def test_incremental_mean_updates_mean_correctly(): - old = np.array([2.0, 2.0]) - new = np.array([4.0, 0.0]) - out = LevelDAG._incremental_mean(old, new, n=2) - np.testing.assert_allclose(out, np.array([3.0, 1.0])) - - def test_reduce_force_and_torque_updates_counts_and_means(): dag = LevelDAG() shared = { - "force_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "torque_covariances": {"ua": {}, "res": [None], "poly": [None]}, - "frame_counts": { + "force_sums": {"ua": {}, "res": [None], "poly": [None]}, + "torque_sums": {"ua": {}, "res": [None], "poly": [None]}, + "force_counts": { + "ua": {}, + "res": np.zeros(1, dtype=int), + "poly": np.zeros(1, dtype=int), + }, + "torque_counts": { "ua": {}, "res": np.zeros(1, dtype=int), "poly": np.zeros(1, dtype=int), }, + "forcetorque_sums": {"res": [None], "poly": [None]}, + "forcetorque_counts": { + "res": np.zeros(1, dtype=int), + "poly": np.zeros(1, dtype=int), + }, + "force_covariances": {}, + "torque_covariances": {}, + "forcetorque_covariances": {}, "group_id_to_index": {9: 0}, } @@ -40,15 +38,24 @@ def test_reduce_force_and_torque_updates_counts_and_means(): frame_out = { "force": {"ua": {(0, 0): F1}, "res": {9: F1}, "poly": {}}, "torque": {"ua": {(0, 0): T1}, "res": {9: T1}, "poly": {}}, + "force_counts": {"ua": {(0, 0): 1}, "res": {9: 1}, "poly": {}}, + "torque_counts": {"ua": {(0, 0): 1}, "res": {9: 1}, "poly": {}}, } dag._reduce_force_and_torque(shared, frame_out) - assert shared["frame_counts"]["ua"][(0, 0)] == 1 + assert shared["force_counts"]["ua"][(0, 0)] == 1 + np.testing.assert_allclose(shared["force_sums"]["ua"][(0, 0)], F1) + np.testing.assert_allclose(shared["torque_sums"]["ua"][(0, 0)], T1) + + assert shared["force_counts"]["res"][0] == 1 + np.testing.assert_allclose(shared["force_sums"]["res"][0], F1) + np.testing.assert_allclose(shared["torque_sums"]["res"][0], T1) + + dag._finalize_means(shared) + np.testing.assert_allclose(shared["force_covariances"]["ua"][(0, 0)], F1) np.testing.assert_allclose(shared["torque_covariances"]["ua"][(0, 0)], T1) - - assert shared["frame_counts"]["res"][0] == 1 np.testing.assert_allclose(shared["force_covariances"]["res"][0], F1) np.testing.assert_allclose(shared["torque_covariances"]["res"][0], T1) @@ -56,7 +63,7 @@ def test_reduce_force_and_torque_updates_counts_and_means(): def test_reduce_forcetorque_no_key_is_noop(): dag = LevelDAG() shared = { - "forcetorque_covariances": {"res": [None], "poly": [None]}, + "forcetorque_sums": {"res": [None], "poly": [None]}, "forcetorque_counts": { "res": np.zeros(1, dtype=int), "poly": np.zeros(1, dtype=int), @@ -66,7 +73,7 @@ def test_reduce_forcetorque_no_key_is_noop(): dag._reduce_forcetorque(shared, frame_out={}) assert shared["forcetorque_counts"]["res"][0] == 0 - assert shared["forcetorque_covariances"]["res"][0] is None + assert shared["forcetorque_sums"]["res"][0] is None def test_run_frame_stage_calls_execute_frame_for_each_ts(simple_ts_list): @@ -81,6 +88,8 @@ def test_run_frame_stage_calls_execute_frame_for_each_ts(simple_ts_list): dag._frame_dag.execute_frame.side_effect = lambda shared_data, frame_index: { "force": {"ua": {}, "res": {}, "poly": {}}, "torque": {"ua": {}, "res": {}, "poly": {}}, + "force_counts": {"ua": {}, "res": {}, "poly": {}}, + "torque_counts": {"ua": {}, "res": {}, "poly": {}}, } dag._reduce_one_frame = MagicMock()