perf(cam): compile the CAM driver — jit + lax.scan (plan 027)

docs/FEATURES.md

@@ -49,7 +49,7 @@ The CAM code line shares leaf kernels with E3SM but none of the orchestration (`
 | coag (`pair_option_acoag = 3`) | `modal_aero_coag.F90:73-990` | **ported (validated)** — three pairs, sequential three-branch number solves, aged-through aitken transfer with shell accounting, coag-side 8-monolayer aging, over the byte-identical `getcoags`. 6e-16 vs `tools/capture_coag`. Found + fixed en route: `shr_const_rgas` is the product 8314.467591, not 8314.46 |
 | rename A1 (`no_acc_crs`) | `modal_aero_rename.F90:243-624` | **ported (validated)** — inside `cam_driver.py`; A2 (`acc_crs`) stays deferred (measured inert below `qso2` ~1e-5); references pinned `no_acc_crs` |
 | calcsize / wateruptake topology threading | shared kernels | **landed** — optional `tables` bundles (E3SM default bit-identical); wateruptake gains `qv=` and the `strat=` wt%-composition water branch |
-| CAM box driver | `mam_box_driver_cam.F90` | **ported (validated end-to-end)** — `cam_run_step`/`cam_run_timesteps`: SO2 stub → calcsize → sulfeq → wateruptake → mmr↔vmr → microphysics, substep loop wrapping the whole step (**default `n_substeps = 16`, ADR-021**; parity tests pin 1). 120 steps × {cam_mam4, cam_mam5} × {trop, strat}: every printed tracer at the reference's 7-digit print floor (~5e-7), total sulfur **4.5e-15**. `tests/test_cam_driver.py` |
+| CAM box driver | `mam_box_driver_cam.F90` | **ported (validated end-to-end)** — `cam_run_step`/`cam_run_timesteps`: SO2 stub → calcsize → sulfeq → wateruptake → mmr↔vmr → microphysics, substep loop wrapping the whole step (**default `n_substeps = 16`, ADR-021**; parity tests pin 1). 120 steps × {cam_mam4, cam_mam5} × {trop, strat}: every printed tracer at the reference's 7-digit print floor (~5e-7), total sulfur **4.5e-15**. `tests/test_cam_driver.py`. **Compiled (plan 027)**: jitted step body + `lax.scan` over substeps and steps; statics per ADR-020 (topology, toggles), rate traced. G5 workload 81.6 s eager → ~1 s cold / **0.02 s warm** |
 
 ## Modes and species
 

docs/PROGRESS.md

@@ -31,6 +31,15 @@ Each entry: date, short title, links to commits / PRs, one-paragraph summary.
 - **Known limitation (deferred)**: the `qaer_del_cond` / `qaer_del_coag` budget-attribution split (F90:5169-5210) is not ported — write-only diagnostics, state evolution unaffected. See `docs/DEFERRED.md`.
 
 ---
+## 2026-08-26 — CAM driver phase-B: jit + lax.scan (branch `feat/cam-driver-jit`)
+
+- PR: stacked on `feat/cam-driver` (#74) while merges to `main` are held. Plan: `docs/plans/027-cam-driver-jit.md`.
+- `cam_run_step`/`cam_run_timesteps` become thin unjitted wrappers (resolving `topology=None` outside the trace) over jitted inners; the substep loop and the step loop are `lax.scan`, with the per-substep reseed flag and the per-step `first_step` flag as traced scan `xs` so both reseed modes share one compilation. Statics per ADR-020: topology, `strat`, `n_substeps`, `do_*`, the two compat flags; `so2_to_h2so4_rate` and the state are traced (rate sweeps don't recompile — locked by test).
+- Scan trajectory bit-identical to repeated single steps; all G5 parity bars unchanged.
+- **Measured** (120 steps × 16 substeps, strat): eager 81.6 s → ~1.0 s cold / **0.02 s warm** (~4000×).
+
+---
+
 
 ## 2026-07-08 — Configurable other-process gas production (`configure_gas_netprod`) (`main`)
 
@@ -42,6 +51,14 @@ Each entry: date, short title, links to commits / PRs, one-paragraph summary.
 - **Tests**: `tests/test_amicphys.py::test_configure_gas_netprod_default_and_override` (defaults, per-rate override, `None` leaves unchanged, save/restore hygiene).
 
 ---
+## 2026-08-27 — Stratospheric RH scenario studies on the CAM driver (branch `feat/cam-driver-jit`)
+
+- First science exploration on the compiled CAM driver (plan 027 made it practical: each 12-h, 16-substep run is ~0.2 s warm). Script: `scripts/cam_rh_scenarios.py`; figures: `docs/figures/cam_strat_rh_{banana,snapshots,budgets}_{6h,12h}.png`. Exploratory, not validation.
+- **6-h strong-forcing sweep** (SO2 = 0.1 ppmv): classic nucleation–growth banana; RH-insensitive number/mass because the regime is production-limited.
+- **12-h weak-forcing experiments** (H2SO4 = 1e7 cm⁻³) isolate the RH → water uptake → wet diameter → condensation-sink chain: RH-ordered gas-depletion floors spanning 3.5 decades (burst, no SO2); ~12% number spread with MORE particles at LOW RH (background SO2 = 1e-10 vmr — the weaker sink leaves gas for nucleation); a 31.5→36.2 nm wet-diameter fan at identical dry size plus 35% lower steady gas at RH 50% (condensation-only). At 232 K, 1e7 cm⁻³ H2SO4 nucleates fully within one 30 s step at any RH — RH-sensitive *number* needs the sustained weak-production regime.
+
+---
+
 
 ## 2026-06-24 — Float32-safe coag + `JAX_ENABLE_X64=0` opt-out (`main`)
 

docs/plans/027-cam-driver-jit.md

@@ -0,0 +1,56 @@
+# Plan 027 — CAM driver phase-B: `jit` + `lax.scan`
+
+**Status:** DONE (2026-08-26; started 2026-08-26, owner: "Can you work on the
+optimization and make the next PR?"). Stacked on `feat/cam-driver`
+(plan 025); PR targets that branch while merges to `main` are held.
+
+## Goal
+
+`cam_run_step` / `cam_run_timesteps` are phase-A eager Python (ADR-004:
+correctness before optimization; plan 025 G4b/G5 established correctness).
+Make them compiled, mirroring what M6 PR-J1/J2 did for the E3SM driver:
+
+1. `cam_run_step` → thin public wrapper (resolves `topology=None` OUTSIDE
+   the trace — `get_topology()` raises inside a jit trace by design) over a
+   jitted inner with the code-path selectors static
+   (`topology`, `strat`, `n_substeps`, `do_*`, `bug_compat_stale_dumfac`,
+   `reseed_dgnwet_each_step`) and the numbers traced
+   (`so2_to_h2so4_rate`, everything in `state` — ADR-020's split).
+   The substep loop becomes `lax.scan`; the per-substep reseed/first-step
+   logic becomes a traced boolean carried as scan `xs`, so both reseed
+   modes share one compiled body.
+2. `cam_run_timesteps` → jitted `lax.scan` over steps (static `n_steps`),
+   stacked trajectory as today. The carry pytree must be stable: the
+   wrapper pre-populates `qaerwat` (zeros) the way the E3SM driver
+   pre-populates calcsize's derived keys.
+
+## Verify
+
+1. Every existing test passes UNCHANGED — the G5 bars (per-tracer 2e-6
+   print floor; totS 1e-13) have orders of headroom over XLA fusion noise.
+   → `pytest tests/test_cam_driver.py` + full suite.
+2. jit-cache semantics locked by tests, per the #65/#77 precedents:
+   a different `so2_to_h2so4_rate` VALUE must NOT retrace (traced leaf);
+   a different topology / `n_substeps` MUST hit a different cache entry
+   (static); scan-trajectory ≡ repeated `cam_run_step` (consistency).
+3. Measured wall-time before/after on the G5 workload (120 steps × 16
+   substeps), recorded in PROGRESS.md.
+
+## Out of scope
+
+`vmap`/sharding (no batched host yet); making `mam_microphysics_cam` and
+below individually jitted (they compile as part of the driver body);
+any physics or API change beyond the wrapper/inner split.
+
+## Results (2026-08-26)
+
+- All existing tests pass unchanged; full G5 parity bars intact.
+- `lax.scan`-over-steps trajectory is **bit-identical** to repeated
+  `cam_run_step` calls (same jitted body, asserted with
+  `assert_array_equal`).
+- Cache semantics locked by `test_rate_is_traced_and_substeps_are_static`:
+  rate sweep reuses one compilation; a new `n_substeps` retraces.
+- **Measured** (G5 workload: 120 steps × 16 substeps, strat, cam_mam4):
+  eager 81.6 s → jitted ~1.0 s cold (compile + run) / **0.02 s warm** —
+  ~4000× warm. The trajectory scan traces the step body once regardless
+  of `n_steps`.