analyze: chunk per-plot aggregation across HTCondor jobs

Add a second parallelism axis to giant analyze: each plot's data can now
be split into a configurable number of event_id-disjoint chunks, each
computed as its own HTCondor job, bounding per-job walltime and scan cost
on large rollout/reference files instead of one job re-scanning the
whole file per plot.

Every PlotSpec now splits into compute_partial (runs per (plot, chunk)
job against a chunk-filtered Bundle) and finalize (merges chunks -
elementwise sum for fixed-edge histograms/species shares, concatenate
-then-recompute for specs that derive edges or mean/std from the full
per-event/per-secondary array). Router diagnostics stay chunkable=False
and always run as a single job. giant analyze render now joins every
plot's chunk partials (merge_all) before rendering, transparently.

New: --chunks on `analyze prep`/`analyze submit`, --chunk on
`analyze compute-one`, and a new `analyze merge-one` command.

tests/test_catalog.py

@@ -2,21 +2,30 @@
 
 from __future__ import annotations
 
+import numpy as np
 import pytest
 
 from giant.analysis import build_catalog, catalog_ids, get_spec
-from giant.analysis.catalog import Bundle
-from giant.analysis.context import build_context
+from giant.analysis.catalog import Bundle, PlotSpec
+from giant.analysis.context import Context, build_context
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
 
-@pytest.fixture(scope="module")
-def bundle() -> Bundle:
+def _build_ctx() -> Context:
     r, t = _rollout_frame(), _reference_frame()
-    ctx = build_context(
+    return build_context(
         r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
     )
-    return Bundle.open(r, t, ctx)
+
+
+@pytest.fixture(scope="module")
+def ctx() -> Context:
+    return _build_ctx()
+
+
+@pytest.fixture(scope="module")
+def bundle(ctx: Context) -> Bundle:
+    return Bundle.open(_rollout_frame(), _reference_frame(), ctx)
 
 
 def test_catalog_ids_unique_and_nonempty():
@@ -36,7 +45,7 @@ def test_get_spec_roundtrip_and_unknown():
 
 def test_every_spec_computes_valid_reduced(bundle: Bundle):
     for spec in build_catalog():
-        r = spec.compute(bundle)
+        r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
         assert r.id == spec.id
         assert r.kind in {
             "overlay_hist",
@@ -81,3 +90,64 @@ def _validate_payload(r) -> None:
             for side in ("rollout", "reference"):
                 if side in p:
                     assert cat in p[side]
+
+
+# ---------------------------------------------------------------------------
+# chunked (compute_partial x N -> finalize) must match the unchunked (N=1) result
+# ---------------------------------------------------------------------------
+
+# One representative id per merge shape: sum-mergeable (marginal_edep,
+# sec_count_per_species via pdg-keyed sums), concat-then-finalize with
+# data-dependent edges (event_total_edep), concat-then-mean/std (shower_
+# longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a
+# ratio (species_edep_share), and a chunkable=False passthrough (router_gating).
+_CHUNK_EQUIVALENCE_IDS = [
+    "marginal_edep",
+    "species_edep_share",
+    "event_total_edep",
+    "shower_longitudinal",
+    "leakage_fraction",
+    "sec_count_per_species",
+    "router_gating",
+]
+
+
+def _assert_payload_close(a, b, path: str = "payload") -> None:
+    """Recursively compare two JSON-shaped payloads (float-tolerant)."""
+    assert type(a) is type(b), f"{path}: {type(a)} != {type(b)}"
+    if isinstance(a, dict):
+        assert set(a) == set(b), f"{path}: key mismatch {set(a)} != {set(b)}"
+        for k in a:
+            _assert_payload_close(a[k], b[k], f"{path}.{k}")
+    elif isinstance(a, list):
+        assert len(a) == len(b), f"{path}: length mismatch"
+        for i, (x, y) in enumerate(zip(a, b)):
+            _assert_payload_close(x, y, f"{path}[{i}]")
+    elif isinstance(a, float):
+        assert np.isclose(a, b, atol=1e-9), f"{path}: {a} != {b}"
+    else:
+        assert a == b, f"{path}: {a} != {b}"
+
+
+@pytest.mark.parametrize("spec_id", _CHUNK_EQUIVALENCE_IDS)
+def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
+    """A plot computed over N event-disjoint chunks then merged must equal the
+    same plot computed in one unchunked pass — the core chunking correctness
+    guarantee (see the analysis-rollout-plots chunking plan)."""
+    spec: PlotSpec = get_spec(spec_id)
+    r, t = _rollout_frame(), _reference_frame()
+
+    unchunked_bundle = Bundle.open(r, t, ctx)
+    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], ctx)
+
+    # 4 chunks over only 2 distinct event_ids also exercises empty chunks.
+    n_chunks = 4 if spec.chunkable else 1
+    parts = [
+        spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks)))
+        for k in range(n_chunks)
+    ]
+    chunked = spec.finalize(parts, ctx)
+
+    assert chunked.id == unchunked.id
+    assert chunked.kind == unchunked.kind
+    _assert_payload_close(unchunked.payload, chunked.payload)

tests/test_condor.py

@@ -16,11 +16,13 @@ from giant.analysis import (
     compute_reduced,
     derive_run_dir,
     load_rollout_yaml,
+    merge_one,
     prep,
     write_submit,
 )
+from giant.analysis.catalog import get_spec
 from giant.analysis.condor import Context
-from giant.analysis.reduced import Reduced
+from giant.analysis.reduced import Partial, Reduced
 from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
@@ -51,11 +53,14 @@ def _write_inputs(tmp_path: Path) -> Path:
     return yaml_path
 
 
-def _prep(rollout_yaml: Path, run_dir: str | Path | None = None) -> Path:
+def _prep(
+    rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1
+) -> Path:
     """``prep`` with small test-sized context bins/sampling."""
     return prep(
         rollout_yaml,
         run_dir,
+        n_chunks=chunks,
         n_energy_bins=2,
         n_marginal_bins=8,
         top_k_pdg=3,
@@ -87,15 +92,16 @@ def test_prep_lays_out_run_dir(tmp_path: Path):
     assert meta.reference.endswith("reference.parquet")
     assert meta.plot_meta["checkpoint"] == "/ckpt/best.pt"
     assert "best.pt" in meta.title
+    assert meta.n_chunks == 1
 
 
 def test_compute_one_from_run_dir(tmp_path: Path):
     run_dir = _prep(_write_inputs(tmp_path))
     out = compute_one("marginal_edep", run_dir)
-    assert out == run_dir / "reduced" / "marginal_edep.json"
-    reduced = Reduced.load(out)
-    assert reduced.id == "marginal_edep"
-    assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1
+    assert out == run_dir / "reduced_partial" / "marginal_edep__0.json"
+    partial = Partial.load(out)
+    assert partial.id == "marginal_edep" and partial.chunk == 0
+    assert "r" in partial.data and "t" in partial.data
 
 
 def test_compute_reduced_explicit_paths(tmp_path: Path):
@@ -108,7 +114,45 @@ def test_compute_reduced_explicit_paths(tmp_path: Path):
         run_dir / "shared.json",
         tmp_path / "r.json",
     )
-    assert Reduced.load(out).id == "marginal_step_length"
+    assert Partial.load(out).id == "marginal_step_length"
+
+
+def test_merge_one_produces_reduced(tmp_path: Path):
+    run_dir = _prep(_write_inputs(tmp_path))
+    compute_one("marginal_edep", run_dir)
+    out = merge_one("marginal_edep", run_dir)
+    assert out == run_dir / "reduced" / "marginal_edep.json"
+    reduced = Reduced.load(out)
+    assert reduced.id == "marginal_edep"
+    assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1
+
+
+def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
+    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    compute_one("marginal_edep", run_dir, chunk_index=0)  # chunk 1 never computed
+    with pytest.raises(FileNotFoundError, match="missing chunk"):
+        merge_one("marginal_edep", run_dir)
+
+
+def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path):
+    (tmp_path / "a").mkdir()
+    (tmp_path / "b").mkdir()
+    unchunked_dir = _prep(_write_inputs(tmp_path / "a"))
+    compute_one("marginal_step_length", unchunked_dir)
+    unchunked = Reduced.load(merge_one("marginal_step_length", unchunked_dir))
+
+    chunked_dir = _prep(_write_inputs(tmp_path / "b"), chunks=2)
+    for k in range(2):
+        compute_one("marginal_step_length", chunked_dir, chunk_index=k)
+    chunked = Reduced.load(merge_one("marginal_step_length", chunked_dir))
+
+    assert chunked.payload == unchunked.payload
+
+
+def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path):
+    run_dir = _prep(_write_inputs(tmp_path))  # n_chunks=1 (default)
+    with pytest.raises(ValueError, match="out of range"):
+        compute_one("marginal_edep", run_dir, chunk_index=1)
 
 
 def test_write_submit_description(tmp_path: Path):
@@ -119,12 +163,15 @@ def test_write_submit_description(tmp_path: Path):
     assert "docker_image = mschnepf/slc7-condocker" in txt
     assert "requirements = TARGET.ProvidesETPResources" in txt
     assert "accounting_group = cms" in txt
-    assert "queue plotid from" in txt
-    assert (run_dir / "plotids.txt").read_text().split() == catalog_ids()
+    assert "queue plotid,chunk from" in txt
+    jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
+    assert [i for i, _ in jobs] == catalog_ids()
+    assert all(k == "0" for _, k in jobs)  # n_chunks=1 default
     wrapper = run_dir / "run_compute.sh"
     assert wrapper.exists() and (wrapper.stat().st_mode & 0o111)
     body = wrapper.read_text()
-    assert "giant analyze compute-one --id" in body and "--run-dir" in body
+    assert "giant analyze compute-one --id" in body
+    assert "--chunk" in body and "--run-dir" in body
 
 
 def test_write_submit_remote_flag(tmp_path: Path):
@@ -135,3 +182,18 @@ def test_write_submit_remote_flag(tmp_path: Path):
     txt = write_submit(cfg).read_text()
     assert "+RemoteJob = True" in txt
     assert "ProvidesETPResources" not in txt
+
+
+def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
+    assert get_spec("router_gating").chunkable is False
+    run_dir = _prep(_write_inputs(tmp_path))
+    cfg = SubmitConfig(
+        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
+    )
+    write_submit(cfg)
+    jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
+    counts: dict[str, int] = {}
+    for spec_id, _ in jobs:
+        counts[spec_id] = counts.get(spec_id, 0) + 1
+    assert counts["marginal_edep"] == 4
+    assert counts["router_gating"] == 1  # chunkable=False, ignores n_chunks