diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py index 914c8ab..13c5183 100644 --- a/giant/analysis/condor.py +++ b/giant/analysis/condor.py @@ -130,7 +130,9 @@ class RunMeta: return cls(**json.loads(Path(path).read_text())) -def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -> list[int]: +def _rows_per_chunk( + rollout: str | Path, reference: str | Path, n_chunks: int +) -> list[int]: """Rollout+reference row count of each ``event_id % n_chunks`` chunk. One cheap streaming ``group_by`` per side (just the ``event_id`` column) — @@ -345,7 +347,9 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> st ) -def _job_walltimes(run_dir: Path, ids: list[str], n_chunks: int) -> list[tuple[str, int, int]]: +def _job_walltimes( + run_dir: Path, ids: list[str], n_chunks: int +) -> list[tuple[str, int, int]]: """``(spec_id, chunk, walltime_s)`` for every job, sized from ``run_meta.json``. Row counts come from ``prep``'s ``RunMeta.rows_per_chunk``/``total_rows``; diff --git a/giant/analysis/runtime_estimate.py b/giant/analysis/runtime_estimate.py index 7f6cc37..f5279f7 100644 --- a/giant/analysis/runtime_estimate.py +++ b/giant/analysis/runtime_estimate.py @@ -54,7 +54,9 @@ _FIXED_OVERHEAD_S = 60.0 # scan. Calibrated from the 3 real router jobs' observed wall times (119, 66, # 124s) — max minus _FIXED_OVERHEAD_S, on top of it. _ROUTER_FIXED_S = 64.0 -_ROUTER_IDS = frozenset({"router_gating", "router_share_by_pdg", "router_share_by_process"}) +_ROUTER_IDS = frozenset( + {"router_gating", "router_share_by_pdg", "router_share_by_process"} +) # Conservative fallback for any catalog id not in _COST_MODEL (e.g. a plot # added after the last calibration run) — the most expensive fitted per-row diff --git a/scripts/profile_analysis_costs.py b/scripts/profile_analysis_costs.py index bbb0a98..de97a2b 100644 --- a/scripts/profile_analysis_costs.py +++ b/scripts/profile_analysis_costs.py @@ -68,9 +68,7 @@ def _make_rollout(n: int, n_events: int, seed: int) -> pl.DataFrame: reasons = np.where( is_synthetic, - rng.choice( - ["escaped", "energy_cutoff", "max_steps", "unknown_pdg"], size=n - ), + rng.choice(["escaped", "energy_cutoff", "max_steps", "unknown_pdg"], size=n), "natural_end", ) @@ -96,7 +94,9 @@ def _make_rollout(n: int, n_events: int, seed: int) -> pl.DataFrame: "post_dx": post_dir[:, 0], "post_dy": post_dir[:, 1], "post_dz": post_dir[:, 2], - "edep": np.where(is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep), + "edep": np.where( + is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep + ), "step_length": np.where(is_synthetic, 0.0, step_length), "material": rng.choice(_MATERIALS, size=n), "layer_id": rng.integers(0, 30, size=n), @@ -163,11 +163,19 @@ def _make_reference(n: int, n_events: int, seed: int) -> pl.DataFrame: ) -def _time(spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path) -> float: +def _time( + spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path +) -> float: t0 = time.perf_counter() compute_reduced( - spec_id, rollout, reference, shared, out, checkpoint=None, - chunk_index=0, n_chunks=1, + spec_id, + rollout, + reference, + shared, + out, + checkpoint=None, + chunk_index=0, + n_chunks=1, ) return time.perf_counter() - t0 @@ -187,8 +195,11 @@ def main() -> None: shared = tmp_path / f"shared_{n_side}.json" ctx = build_context( - rollout, reference, - n_energy_bins=4, n_marginal_bins=50, top_k_pdg=6, + rollout, + reference, + n_energy_bins=4, + n_marginal_bins=50, + top_k_pdg=6, sample_rows=min(n_side, 200_000), ) ctx.save(shared)