diff --git a/docs/FEATURES.md b/docs/FEATURES.md index 3982eec..6534172 100644 --- a/docs/FEATURES.md +++ b/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 diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 05689dc..24d2f59 100644 --- a/docs/PROGRESS.md +++ b/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`) diff --git a/docs/figures/cam_strat_rh_banana_12h.png b/docs/figures/cam_strat_rh_banana_12h.png new file mode 100644 index 0000000..de9caeb Binary files /dev/null and b/docs/figures/cam_strat_rh_banana_12h.png differ diff --git a/docs/figures/cam_strat_rh_banana_6h.png b/docs/figures/cam_strat_rh_banana_6h.png new file mode 100644 index 0000000..c76b61a Binary files /dev/null and b/docs/figures/cam_strat_rh_banana_6h.png differ diff --git a/docs/figures/cam_strat_rh_budgets_12h.png b/docs/figures/cam_strat_rh_budgets_12h.png new file mode 100644 index 0000000..d7bd023 Binary files /dev/null and b/docs/figures/cam_strat_rh_budgets_12h.png differ diff --git a/docs/figures/cam_strat_rh_budgets_6h.png b/docs/figures/cam_strat_rh_budgets_6h.png new file mode 100644 index 0000000..7d7aa9d Binary files /dev/null and b/docs/figures/cam_strat_rh_budgets_6h.png differ diff --git a/docs/figures/cam_strat_rh_snapshots_12h.png b/docs/figures/cam_strat_rh_snapshots_12h.png new file mode 100644 index 0000000..7a1eec2 Binary files /dev/null and b/docs/figures/cam_strat_rh_snapshots_12h.png differ diff --git a/docs/figures/cam_strat_rh_snapshots_6h.png b/docs/figures/cam_strat_rh_snapshots_6h.png new file mode 100644 index 0000000..a50bb1b Binary files /dev/null and b/docs/figures/cam_strat_rh_snapshots_6h.png differ diff --git a/docs/plans/027-cam-driver-jit.md b/docs/plans/027-cam-driver-jit.md new file mode 100644 index 0000000..c267bed --- /dev/null +++ b/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`. diff --git a/mam4_jax/coupling/cam_driver.py b/mam4_jax/coupling/cam_driver.py index 0857a96..36264b9 100644 --- a/mam4_jax/coupling/cam_driver.py +++ b/mam4_jax/coupling/cam_driver.py @@ -53,6 +53,7 @@ import functools +import jax import jax.numpy as jnp import numpy as np @@ -927,58 +928,34 @@ def _cam_wateruptake_tables(tb: _CamTables): rhcrystal=topo.rhcrystal_amode, rhdeliques=topo.rhdeliques_amode, ) - -def cam_run_step(state, *, topology=None, strat=False, n_substeps=16, - do_calcsize=True, do_wateruptake=True, - do_gasaerexch=True, do_newnuc=True, do_coag=True, - so2_to_h2so4_rate=1.0e-5, - bug_compat_stale_dumfac=False, - first_step=False, - reseed_dgnwet_each_step=True): - """One CAM box-model step (mam_box_run_cam's loop body). - - Per SUBSTEP of ``deltat / n_substeps`` — sub-stepping wraps the WHOLE - physics, so ``n_substeps = n`` is semantically identical to running - the box at ``deltat/n``, the exact quantity the Fortran - dt-convergence study varied. - - **``n_substeps`` defaults to 16 (ADR-021)** — a deliberate deviation - from the defaults-reproduce-the-reference convention, owner-decided - (plan 025 A6): CAM's un-substepped sequential splitting does not - converge while nucleation is active, and n = 1 (CAM's own behaviour, - and the reference box's) is 26-78%% from the converged answer at - dt = 30 s on the default scenario. The error is first-order (halves - per doubling); 16 lands at ~1.4-4.4%% for 16x cost. **Pass - ``n_substeps=1`` to reproduce CAM / the Fortran reference exactly** - — every parity test does. The full measured table is in ADR-021 and - plan 025 §7. - - The per-substep sequence: - - 1. SO2 → H2SO4 first-order conversion (the gas-chemistry stub, - sulfur-mole-conserving by MW ratio; default rate 1e-5 /s = - ``nl_so2_to_h2so4_rate``), producing ``del_h2so4_gasprod`` as a - vmr increment for nucleation; - 2. calcsize (mmr, aitacc transfer on, the fixed per-mode ``dumfac`` - by default — ``bug_compat_stale_dumfac=True`` reproduces upstream - CAM including the bug, matching the non-fixdumfac reference - builds); - 3. ``sulfeq`` when ``strat``: per-mode equilibrium H2SO4 from the - Tabazadeh cluster at the LAGGED wet diameter — the incoming - ``dgncur_awet``, i.e. the previous substep's, exactly as CAM's - pbuf carries it (plan 024 §6's explicit carried state); - 4. wateruptake (mmr; ``qv`` passed explicitly; updates - ``dgncur_awet``/``qaerwat``/``wetdens``); - 5. mmr → vmr over the gas window (``q·mwdry/adv_mass``, number - tracers included with their mechanism ``adv_mass``, exactly as - the box does), the microphysics chain, vmr → mmr back. - - STATE (a dict; all tracers MASS mixing ratio over the gas window): - ``q`` (..., gas_pcnst); ``qv``, ``t``, ``pmid``, ``zm``, ``pblh`` - (...,); ``dgncur_a``, ``dgncur_awet``, ``wetdens`` (..., nmodes); - ``deltat`` scalar. ``qqcw`` is identically zero in scope and not - carried. Returns the updated state dict (same keys, plus - ``qaerwat`` once wateruptake has run). +_STEP_STATICS = ("topology", "strat", "n_substeps", "do_calcsize", + "do_wateruptake", "do_gasaerexch", "do_newnuc", "do_coag", + "bug_compat_stale_dumfac", "reseed_dgnwet_each_step") + +#: evolving per-step state; everything else in the state dict is constant +#: forcing echoed through the scan carry untouched. +_EVOLVING = ("q", "dgncur_a", "dgncur_awet", "wetdens", "qaerwat") + + +@functools.partial(jax.jit, static_argnames=_STEP_STATICS) +def _cam_run_step_jit(state, so2_to_h2so4_rate, first_step, *, + topology, strat, n_substeps, + do_calcsize, do_wateruptake, + do_gasaerexch, do_newnuc, do_coag, + bug_compat_stale_dumfac, reseed_dgnwet_each_step): + """Jitted one-step body — plan 027. See :func:`cam_run_step`. + + Code-path selectors are STATIC (each combination compiles once, per + the ADR-020 split); ``so2_to_h2so4_rate`` and ``first_step`` are + traced operands, so sweeping the rate or flipping the first-step + flag reuses one compilation. The substep loop is a ``lax.scan`` + whose ``xs`` is the per-substep reseed flag — both reseed modes + (the box shim's every-substep, production CAM's first-only) share + one compiled body. + + ``topology`` must be a concrete ``Topology`` here (the PUBLIC + wrapper resolves ``None`` before the trace — ``get_topology()`` + deliberately raises inside a jit trace, core/topology.py). """ from mam4_jax.physics.calcsize import calcsize from mam4_jax.physics.strat_sulfate import calc_h2so4_equilib_mixrat @@ -990,13 +967,8 @@ def cam_run_step(state, *, topology=None, strat=False, n_substeps=16, adv = jnp.asarray(tb.adv_mass) mwdry = tb.mwdry dt_s = state["deltat"] / n_substeps - - q = jnp.asarray(state["q"], dtype=jnp.float64) - dgncur_a = jnp.asarray(state["dgncur_a"], dtype=jnp.float64) - dgncur_awet = jnp.asarray(state["dgncur_awet"], dtype=jnp.float64) - wetdens = jnp.asarray(state["wetdens"], dtype=jnp.float64) - qaerwat = state.get("qaerwat") qv, t, pmid = state["qv"], state["t"], state["pmid"] + zm, pblh = state["zm"], state["pblh"] so4_slot = np.full(tb.topology.nmodes, -1, dtype=int) so4_type = tb.topology.specname_amode.index("so4") @@ -1006,10 +978,18 @@ def cam_run_step(state, *, topology=None, strat=False, n_substeps=16, so4_slot[m] = s_ dens_so4 = tb.topology.specdens_amode[so4_type] - for isub in range(n_substeps): + if reseed_dgnwet_each_step: + sub_reseed = jnp.ones(n_substeps, dtype=bool) + else: + sub_reseed = jnp.zeros(n_substeps, dtype=bool).at[0].set(first_step) + + def _substep(carry, reseed): + q, dgncur_a, dgncur_awet, wetdens, qaerwat = carry + # 1. SO2 -> H2SO4 stub (mam_box_driver_cam.F90:363-373). if tb.l_so2 >= 0: - prod = q[..., tb.l_so2] * (-jnp.expm1(-so2_to_h2so4_rate * dt_s)) + prod = q[..., tb.l_so2] * ( + -jnp.expm1(-so2_to_h2so4_rate * dt_s)) q = q.at[..., tb.l_so2].add(-prod) q = q.at[..., tb.l_h2so4].add( prod * (tb.adv_mass[tb.l_h2so4] / tb.adv_mass[tb.l_so2])) @@ -1028,26 +1008,17 @@ def cam_run_step(state, *, topology=None, strat=False, n_substeps=16, # CAM's wateruptake_dr seeds the lagged wet diameter from the # CURRENT dry one under is_first_step() # (modal_aero_wateruptake.F90:329-331) — BEFORE the sulfeq loop - # reads it. In production CAM that fires once (step 0), making - # sulfeq a genuinely LAGGED-wet-diameter quantity (plan 024 §6's - # carried state). In the box REFERENCE, the time-manager shim's - # is_first_step() is TRUE EVERY STEP (vendor - # time_manager.F90:13-22, never advanced by the driver), so the - # reference recomputes sulfeq from the fresh post-calcsize DRY - # diameters each step and the lag never survives. - # ``reseed_dgnwet_each_step=True`` (default) reproduces the - # reference; False gives the production-CAM lagged behaviour - # (seed on the first step only). Without any seed, step 1's - # sulfeq would see dgncur_awet = 0 and the Kelvin clamp - # evaporates every so4 mode (measured: so4_a1 off by 1e4x). - if do_wateruptake and (reseed_dgnwet_each_step - or (first_step and isub == 0)): - dgncur_awet = dgncur_a - - # 3. sulfeq + Tabazadeh composition at the LAGGED wet diameter - # (before wateruptake updates it) — modal_aero_wateruptake.F90:392-396. + # reads it. The box reference's time-manager shim makes that + # true EVERY step (plan 025 §7 G5 finding 1); production CAM + # lags genuinely. ``reseed`` is the traced per-substep flag. + if do_wateruptake: + dgncur_awet = jnp.where(reseed, dgncur_a, dgncur_awet) + + # 3. sulfeq + Tabazadeh composition at the (possibly reseeded) + # pre-wateruptake wet diameter (wateruptake.F90:392-396). if strat: - dmean = dgncur_awet * jnp.exp(1.5 * jnp.asarray(tb.alnsg) ** 2) + dmean = dgncur_awet * jnp.exp( + 1.5 * jnp.asarray(tb.alnsg) ** 2) sulfeq, wtpct, sulden = calc_h2so4_equilib_mixrat( jnp.asarray(t)[..., None], jnp.asarray(pmid)[..., None], jnp.asarray(qv)[..., None], dmean) @@ -1057,10 +1028,7 @@ def cam_run_step(state, *, topology=None, strat=False, n_substeps=16, sulfeq = None strat_wu = None - # 4. wateruptake (mmr; qv explicit; the strat branch replaces - # Köhler with the wt%-composition solution volume — the box sets - # its tropopause above the single level, so under strat that - # branch is live everywhere). + # 4. wateruptake (mmr; qv explicit; strat branch per plan 025). if do_wateruptake: wu = wateruptake( {"q": q, "dgncur_a": dgncur_a, "t": t, "pmid": pmid, @@ -1073,32 +1041,159 @@ def cam_run_step(state, *, topology=None, strat=False, n_substeps=16, # 5. the microphysics chain, on vmr. q_vmr = q * (mwdry / adv) q_vmr = mam_microphysics_cam( - q_vmr, t, pmid, dt_s, qv, state["zm"], state["pblh"], + q_vmr, t, pmid, dt_s, qv, zm, pblh, dgncur_a, dgncur_awet, wetdens, del_h2so4_gasprod, topology=topology, sulfeq=sulfeq, do_gasaerexch=do_gasaerexch, do_newnuc=do_newnuc, do_coag=do_coag) q = q_vmr * (adv / mwdry) - out = {**state, "q": q, "dgncur_a": dgncur_a, - "dgncur_awet": dgncur_awet, "wetdens": wetdens} - if qaerwat is not None: - out["qaerwat"] = qaerwat - return out - - -def cam_run_timesteps(state, n_steps, **kwargs): - """Run ``n_steps`` CAM box steps; return the final state plus a - trajectory dict of stacked per-step snapshots of ``q``, - ``dgncur_a``, ``dgncur_awet``, ``wetdens`` (post-step values, - matching the Fortran box's per-step output rows). Plain Python loop - — phase A; jit/scan is the phase-B optimization pass.""" - import jax - - traj = {k: [] for k in ("q", "dgncur_a", "dgncur_awet", "wetdens")} - for istep in range(n_steps): - state = cam_run_step(state, first_step=(istep == 0), **kwargs) - for k in traj: - traj[k].append(state[k]) - stacked = {k: jax.numpy.stack(v) for k, v in traj.items()} - return state, stacked + return (q, dgncur_a, dgncur_awet, wetdens, qaerwat), None + + carry0 = tuple(jnp.asarray(state[k], dtype=jnp.float64) + for k in _EVOLVING) + carry, _ = jax.lax.scan(_substep, carry0, sub_reseed) + return {**state, **dict(zip(_EVOLVING, carry))} + + +@functools.partial(jax.jit, static_argnames=("n_steps",) + _STEP_STATICS) +def _cam_run_timesteps_jit(state, so2_to_h2so4_rate, *, + n_steps, topology, strat, n_substeps, + do_calcsize, do_wateruptake, + do_gasaerexch, do_newnuc, do_coag, + bug_compat_stale_dumfac, + reseed_dgnwet_each_step): + """Jitted trajectory — ``lax.scan`` over steps (static ``n_steps``), + each step the jitted step body (inlined by the outer trace). The + per-step ``first_step`` flag is the scan ``xs``.""" + statics = dict( + topology=topology, strat=strat, n_substeps=n_substeps, + do_calcsize=do_calcsize, do_wateruptake=do_wateruptake, + do_gasaerexch=do_gasaerexch, do_newnuc=do_newnuc, + do_coag=do_coag, bug_compat_stale_dumfac=bug_compat_stale_dumfac, + reseed_dgnwet_each_step=reseed_dgnwet_each_step) + + def _body(carry_state, first): + new = _cam_run_step_jit(carry_state, so2_to_h2so4_rate, first, + **statics) + out = {k: new[k] for k in ("q", "dgncur_a", "dgncur_awet", + "wetdens")} + return new, out + + firsts = jnp.arange(n_steps) == 0 + return jax.lax.scan(_body, state, firsts) + + +def _prepare_state(state): + """Stable-pytree precondition for the scan carry: ``qaerwat`` exists + (zeros before the first wateruptake, exactly the pbuf's physpkg + default) — same pattern as the E3SM driver's placeholder keys.""" + if "qaerwat" not in state: + state = {**state, + "qaerwat": jnp.zeros_like( + jnp.asarray(state["dgncur_awet"]))} + return state + + +def cam_run_step(state, *, topology=None, strat=False, n_substeps=16, + do_calcsize=True, do_wateruptake=True, + do_gasaerexch=True, do_newnuc=True, do_coag=True, + so2_to_h2so4_rate=1.0e-5, + bug_compat_stale_dumfac=False, + first_step=False, + reseed_dgnwet_each_step=True): + """One CAM box-model step (mam_box_run_cam's loop body), compiled. + + Per SUBSTEP of ``deltat / n_substeps`` — sub-stepping wraps the WHOLE + physics, so ``n_substeps = n`` is semantically identical to running + the box at ``deltat/n``, the exact quantity the Fortran + dt-convergence study varied. + + **``n_substeps`` defaults to 16 (ADR-021)** — a deliberate deviation + from the defaults-reproduce-the-reference convention, owner-decided + (plan 025 A6): CAM's un-substepped sequential splitting does not + converge while nucleation is active, and n = 1 (CAM's own behaviour, + and the reference box's) is 26-78% from the converged answer at + dt = 30 s on the default scenario. The error is first-order (halves + per doubling); 16 lands at ~1.4-4.4% for 16x cost. **Pass + ``n_substeps=1`` to reproduce CAM / the Fortran reference exactly** + — every parity test does. The full measured table is in ADR-021 and + plan 025 §7. + + The per-substep sequence: + + 1. SO2 → H2SO4 first-order conversion (the gas-chemistry stub, + sulfur-mole-conserving by MW ratio; default rate 1e-5 /s = + ``nl_so2_to_h2so4_rate``), producing ``del_h2so4_gasprod`` as a + vmr increment for nucleation; + 2. calcsize (mmr, aitacc transfer on, the fixed per-mode ``dumfac`` + by default — ``bug_compat_stale_dumfac=True`` reproduces upstream + CAM including the bug, matching the non-fixdumfac reference + builds); + 3. ``sulfeq`` when ``strat``: per-mode equilibrium H2SO4 from the + Tabazadeh cluster at the pre-wateruptake wet diameter — + re-seeded from the current dry diameter per + ``reseed_dgnwet_each_step`` (True = the box reference's + behaviour, its ``is_first_step()`` shim being true every step; + False = production CAM's genuine lag, seeded on + ``first_step`` only); + 4. wateruptake (mmr; ``qv`` passed explicitly; updates + ``dgncur_awet``/``qaerwat``/``wetdens``); + 5. mmr → vmr over the gas window (``q·mwdry/adv_mass``, number + tracers included with their mechanism ``adv_mass``, exactly as + the box does), the microphysics chain, vmr → mmr back. + + STATE (a dict; all tracers MASS mixing ratio over the gas window): + ``q`` (..., gas_pcnst); ``qv``, ``t``, ``pmid``, ``zm``, ``pblh`` + (...,); ``dgncur_a``, ``dgncur_awet``, ``wetdens`` (..., nmodes); + ``deltat`` scalar. ``qqcw`` is identically zero in scope and not + carried. Returns the updated state dict (same keys, plus + ``qaerwat``). + + **Compiled (plan 027).** This wrapper resolves ``topology=None`` to + the active topology OUTSIDE the trace and hands everything to a + jitted inner whose code-path selectors (topology, ``strat``, + ``n_substeps``, the ``do_*`` toggles, the two compat flags) are + STATIC — one compilation per combination — while + ``so2_to_h2so4_rate``, ``first_step`` and the state are traced + (sweeping the rate does not recompile; ADR-020's split). + """ + if topology is None: + topology = get_topology() + return _cam_run_step_jit( + _prepare_state(state), so2_to_h2so4_rate, jnp.asarray(first_step), + topology=topology, strat=strat, n_substeps=n_substeps, + do_calcsize=do_calcsize, do_wateruptake=do_wateruptake, + do_gasaerexch=do_gasaerexch, do_newnuc=do_newnuc, do_coag=do_coag, + bug_compat_stale_dumfac=bug_compat_stale_dumfac, + reseed_dgnwet_each_step=reseed_dgnwet_each_step) + + +def cam_run_timesteps(state, n_steps, *, topology=None, strat=False, + n_substeps=16, + do_calcsize=True, do_wateruptake=True, + do_gasaerexch=True, do_newnuc=True, do_coag=True, + so2_to_h2so4_rate=1.0e-5, + bug_compat_stale_dumfac=False, + reseed_dgnwet_each_step=True): + """Run ``n_steps`` CAM box steps; return ``(final_state, trajectory)`` + with per-step post-step snapshots of ``q``, ``dgncur_a``, + ``dgncur_awet``, ``wetdens`` stacked on a leading axis (matching the + Fortran box's output rows). + + **Compiled (plan 027)**: a jitted ``lax.scan`` over steps with + ``n_steps`` static (one cache entry per distinct count — the M6 + PR-J2 pattern), the step body inlined once. The first step's + ``first_step`` flag rides the scan ``xs``, so both reseed modes + share the compilation. Keyword arguments as :func:`cam_run_step`. + """ + if topology is None: + topology = get_topology() + return _cam_run_timesteps_jit( + _prepare_state(state), so2_to_h2so4_rate, + n_steps=n_steps, topology=topology, strat=strat, + n_substeps=n_substeps, + do_calcsize=do_calcsize, do_wateruptake=do_wateruptake, + do_gasaerexch=do_gasaerexch, do_newnuc=do_newnuc, do_coag=do_coag, + bug_compat_stale_dumfac=bug_compat_stale_dumfac, + reseed_dgnwet_each_step=reseed_dgnwet_each_step) diff --git a/scripts/cam_rh_scenarios.py b/scripts/cam_rh_scenarios.py new file mode 100644 index 0000000..139ffd6 --- /dev/null +++ b/scripts/cam_rh_scenarios.py @@ -0,0 +1,357 @@ +"""Stratospheric RH scenario studies on the CAM driver. + +Regenerates ``docs/figures/cam_strat_rh_*.png``. Two studies, both +``cam_mam5``, ``strat=True``, T=232 K, p=50 hPa, dt=30 s, and the ADR-021 +default ``n_substeps=16``: + +* **6-h sweep, strong forcing** (SO2 = 1e-7 vmr = 0.1 ppmv, volcanic-plume + scale): the classic nucleation–growth banana. Nearly RH-insensitive in + number and mass because the system is PRODUCTION-limited — every H2SO4 + molecule produced (1.6e6 cm^-3 s^-1) is consumed at any RH. +* **12-h sweep × three experiments** at 1e7 molec/cm3 initial H2SO4, + isolating the RH → water uptake → wet diameter → condensation-sink + chain: (1) *burst* (no SO2) — RH-ordered gas-depletion floors spanning + 3.5 decades; (2) *background* (SO2 = 1e-10 vmr) — a ~12% number spread + with MORE particles at LOW RH (the weaker sink leaves gas standing for + nucleation); (3) *condensation-only* (nucleation off) — the wet aitken + diameter fans 31.5 → 36.2 nm at identical dry size, and the steady gas + sits 35% lower at RH 50%. + +Exploratory scenario study, not validation — no Fortran reference is +attached. Figure conventions: per-mode categorical colors in fixed order +across every figure; ordered dimensions (time, RH) use sequential ramps. + +Usage: python scripts/cam_rh_scenarios.py [outdir] (default docs/figures) +Runtime: ~2 min on CPU (compile-dominated; the runs themselves are ~0.2 s +each on the jitted driver). +""" +from __future__ import annotations + +import sys +import tempfile +from pathlib import Path + +import numpy as np +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.colors import LogNorm # noqa: E402 + +import jax.numpy as jnp # noqa: E402 + +import mam4_jax # noqa: F401,E402 (enables x64) +import mam4_jax.coupling.cam_driver as cd # noqa: E402 +from mam4_jax.core.cam_params import CAM_PARAMS # noqa: E402 +from mam4_jax.core.cam_topologies import CAM_MAM5 # noqa: E402 +from mam4_jax.physics.cam_saturation import qsat_cam # noqa: E402 + +TOPO = CAM_MAM5 +T, P, DT = 232.0, 5.0e3, 30.0 +N_AIR = P / (1.380649e-23 * T) * 1e-6 # molec/cm3 +RHO = P / (287.0423 * T) # kg/m3 +RHS = [0.05, 0.10, 0.20, 0.50] +RH_TAGS = [f"rh{int(r * 100):02d}" for r in RHS] +RH_LBL = {t: f"RH {int(t[2:])}%" for t in RH_TAGS} +MODES = ["accum", "aitken", "coarse", "prim. carbon", "coarse_strat"] +MODE_C = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4"] # fixed order +RH_C = ["#b8d3f2", "#7fb0e6", "#4a90dc", "#1c5cab"] # sequential +GREYS = ["#c9c9c9", "#a0a0a0", "#747474", "#424242", "#000000"] +NUMC = (1.0e8, 1.0e9, 1.0e5, 1.0e0) # #/m3 per mode 1..4 (box namelist) +NUMC5 = 1.0e4 +SO4FRAC = (1.0, 1.0, 1.0, 0.0) +GRID = dict(color="#e3e3e3", lw=0.6) + +_tb = cd._cam_tables(TOPO) +_names = CAM_PARAMS[TOPO.name]["cnst_names"] +_iH = _names.index("H2SO4") +SIG = np.asarray(TOPO.sigmag_amode) +LNS = np.log(SIG) +DP = np.logspace(np.log10(1e-9), np.log10(3e-6), 240) +LNDP = np.log(DP) + + +def build_ic(rh, qso2, qh2so4): + """The box driver's IC construction (mam_box_driver_cam.F90:279-345).""" + p = CAM_PARAMS[TOPO.name] + q = np.zeros(p["gas_pcnst"]) + for m in range(TOPO.nmodes): + numkg = (NUMC[m] if m < 4 else NUMC5) / RHO + q[_tb.num_ptr[m]] = numkg + dg, sg = TOPO.dgnum_amode[m], TOPO.sigmag_amode[m] + tmpvol = numkg * (np.pi / 6) * dg ** 3 * np.exp(4.5 * np.log(sg) ** 2) + if _tb.lptr_so4[m] >= 0: + frac = 1.0 if m == 4 else SO4FRAC[m] + q[_tb.lptr_so4[m]] = frac * tmpvol * 1770.0 + q[_iH] = qh2so4 * p["adv_mass"][_iH] / p["mwdry"] + q[_names.index("SO2")] = qso2 * p["adv_mass"][_names.index("SO2")] / p["mwdry"] + _es, qs = qsat_cam(jnp.asarray(T), jnp.asarray(P)) + nm = TOPO.nmodes + return {"q": jnp.asarray(q), "qv": jnp.asarray(rh * float(qs)), + "t": jnp.asarray(T), "pmid": jnp.asarray(P), + "zm": jnp.asarray(500.0), "pblh": jnp.asarray(1000.0), + "dgncur_a": jnp.zeros(nm), "dgncur_awet": jnp.zeros(nm), + "wetdens": jnp.zeros(nm), "deltat": jnp.asarray(DT)} + + +def run(rh, nstep, qso2, qh2so4, do_newnuc=True): + _, traj = cd.cam_run_timesteps(build_ic(rh, qso2, qh2so4), nstep, + topology=TOPO, strat=True, n_substeps=16, + do_newnuc=do_newnuc) + q = np.asarray(traj["q"]) + return dict( + num=q[:, _tb.num_ptr], + so4=np.stack([q[:, _tb.lptr_so4[m]] if _tb.lptr_so4[m] >= 0 + else np.zeros(nstep) for m in range(TOPO.nmodes)], 1), + dgn=np.asarray(traj["dgncur_a"]), + dgnwet=np.asarray(traj["dgncur_awet"]), + h2so4=q[:, _iH] * _tb.mwdry / _tb.adv_mass[_iH] * N_AIR, # molec/cm3 + ) + + +def init_dist_inputs(): + num0 = np.asarray([(NUMC[m] if m < 4 else NUMC5) / RHO + for m in range(TOPO.nmodes)]) + return num0, np.asarray(TOPO.dgnum_amode) + + +def dist(num, dgn): + """dN/dlnD (#/mg-air) on the DP grid; also per-mode components.""" + tot = np.zeros_like(DP) + pm = [] + for m in range(len(num)): + c = (num[m] * 1e-6) / (np.sqrt(2 * np.pi) * LNS[m]) * np.exp( + -((LNDP - np.log(dgn[m])) ** 2) / (2 * LNS[m] ** 2)) + pm.append(c) + tot += c + return tot, pm + + +def _style(): + plt.rcParams.update({"font.size": 9, "axes.titlesize": 9.5, + "axes.edgecolor": "#999999", "axes.linewidth": 0.8}) + + +def banana_panel(ax, run_d, t_hr, first): + num, dgn = run_d["num"], run_d["dgn"] + Z = np.zeros((len(DP), len(t_hr))) + for it in range(len(t_hr)): + Z[:, it], _ = dist(num[it], dgn[it]) + pm = ax.pcolormesh(t_hr, DP * 1e9, np.maximum(Z, 1e-2), + norm=LogNorm(vmin=1e0, vmax=None), + cmap="Blues", rasterized=True, shading="auto") + for m in range(TOPO.nmodes): + if num[:, m].max() * 1e-6 > 1e-3: + ax.plot(t_hr, dgn[:, m] * 1e9, color=MODE_C[m], lw=1.4, + label=MODES[m] if first else None) + ax.set_yscale("log") + ax.set_ylim(1, 3000) + return pm + + +def fig_sweep_6h(out: Path): + """Strong-forcing 6-h sweep: banana + snapshots + budgets.""" + nstep = 720 + t_hr = (np.arange(nstep) + 1) * DT / 3600.0 + runs = {tag: run(rh, nstep, qso2=1.0e-7, qh2so4=1.0e-11) + for tag, rh in zip(RH_TAGS, RHS)} + + fig, axes = plt.subplots(2, 2, figsize=(11.5, 7.6), sharex=True, sharey=True) + for i, (ax, tag) in enumerate(zip(axes.ravel(), RH_TAGS)): + pm = banana_panel(ax, runs[tag], t_hr, first=(i == 0)) + pm.set_norm(LogNorm(vmin=1e2, vmax=5e6)) + ax.set_title(RH_LBL[tag], loc="left") + for ax in axes[1]: + ax.set_xlabel("time (h)") + for ax in axes[:, 0]: + ax.set_ylabel("dry diameter $D_p$ (nm)") + fig.colorbar(pm, ax=axes, label=r"dN/dln$D_p$ (# mg$^{-1}$)", pad=0.02) + fig.legend(*axes[0, 0].get_legend_handles_labels(), loc="lower center", + ncol=5, frameon=False, bbox_to_anchor=(0.45, -0.035), + title="mode $d_{gn}$ overlays") + fig.suptitle("Stratospheric nucleation–growth (banana) — CAM driver, cam_mam5, " + "T=232 K, p=50 hPa, SO$_2$=1e-7 vmr", y=0.995, fontsize=11) + fig.savefig(out / "cam_strat_rh_banana_6h.png", dpi=160, bbox_inches="tight") + plt.close(fig) + + num0, dgn0 = init_dist_inputs() + snap_h = [0.0, 0.5, 1.0, 3.0, 6.0] + fig, axes = plt.subplots(2, 2, figsize=(11, 7.6), sharex=True, sharey=True) + for i, (ax, tag) in enumerate(zip(axes.ravel(), RH_TAGS)): + num, dgn = runs[tag]["num"], runs[tag]["dgn"] + for h, c in zip(snap_h, GREYS): + if h == 0.0: + tot, _ = dist(num0, dgn0) + lbl = "initial" + else: + it = int(h * 3600 / DT) - 1 + tot, _ = dist(num[it], dgn[it]) + lbl = f"{h:g} h" + ax.plot(DP * 1e9, tot, color=c, lw=1.7, label=lbl) + _, pm6 = dist(num[-1], dgn[-1]) + for m in range(TOPO.nmodes): + if num[-1, m] > 1e-3: + ax.plot(DP * 1e9, pm6[m], color=MODE_C[m], lw=1.1, ls="--", + label=MODES[m] if i == 0 else None) + ax.set_xscale("log"); ax.set_yscale("log") + ax.set_xlim(1, 3000); ax.set_ylim(1e-2, 3e7) + ax.grid(True, which="major", **GRID) + ax.set_title(RH_LBL[tag], loc="left") + for ax in axes[1]: + ax.set_xlabel("dry diameter $D_p$ (nm)") + for ax in axes[:, 0]: + ax.set_ylabel(r"dN/dln$D_p$ (# mg$^{-1}$)") + h_, l_ = axes[0, 0].get_legend_handles_labels() + axes[0, 0].legend(h_, l_, loc="upper right", fontsize=7.6, ncol=2, + framealpha=0.95) + fig.suptitle("Number-distribution snapshots (grey = time; dashed = per-mode at 6 h)", + y=0.995, fontsize=11) + fig.savefig(out / "cam_strat_rh_snapshots_6h.png", dpi=160, bbox_inches="tight") + plt.close(fig) + + fig, axes = plt.subplots(2, 2, figsize=(11, 7.2)) + axN, axM, ax5, ax50 = axes.ravel() + for tag, c in zip(RH_TAGS, RH_C): + axN.plot(t_hr, runs[tag]["num"].sum(1) * 1e-6, color=c, lw=1.8, + label=RH_LBL[tag]) + axM.plot(t_hr, runs[tag]["so4"].sum(1) * 1e9, color=c, lw=1.8) + axN.set_yscale("log"); axN.set_ylabel(r"total number (# mg$^{-1}$)") + axN.set_title("(a) total number concentration", loc="left") + axM.set_ylabel(r"total SO$_4$ aerosol mass (µg kg$^{-1}$)") + axM.set_title("(b) total sulfate aerosol mass", loc="left") + axN.legend(fontsize=8, framealpha=0.95) + for ax, tag in ((ax5, "rh05"), (ax50, "rh50")): + for m in range(TOPO.nmodes): + if runs[tag]["num"][:, m].max() * 1e-6 > 1e-3: + ax.plot(t_hr, runs[tag]["num"][:, m] * 1e-6, color=MODE_C[m], + lw=1.5, label=MODES[m]) + ax.set_yscale("log"); ax.set_ylabel(r"mode number (# mg$^{-1}$)") + ax5.set_title("(c) per-mode number — RH 5%", loc="left") + ax50.set_title("(d) per-mode number — RH 50%", loc="left") + ax5.legend(fontsize=8, framealpha=0.95) + for ax in axes.ravel(): + ax.grid(True, **GRID); ax.set_xlabel("time (h)") + fig.suptitle("Number and mass budgets — stratospheric RH sweep", y=0.998, + fontsize=11) + fig.tight_layout() + fig.savefig(out / "cam_strat_rh_budgets_6h.png", dpi=160, bbox_inches="tight") + plt.close(fig) + + +def fig_experiments_12h(out: Path): + """Weak-forcing 12-h experiments: burst / background / condensation-only.""" + nstep = 1440 + t_hr = (np.arange(nstep) + 1) * DT / 3600.0 + vmr_1e7 = 1.0e7 / N_AIR + exps = {"burst": (0.0, vmr_1e7, True), + "background": (1.0e-10, vmr_1e7, True), + "condonly": (1.0e-10, vmr_1e7, False)} + exp_lbl = {"burst": "1) burst: H$_2$SO$_4$=1e7 cm$^{-3}$, no SO$_2$", + "background": "2) background: SO$_2$=1e-10 vmr", + "condonly": "3) condensation-only, SO$_2$=1e-10"} + runs = {(e, tag): run(rh, nstep, qso2=s, qh2so4=h, do_newnuc=nuc) + for e, (s, h, nuc) in exps.items() + for tag, rh in zip(RH_TAGS, RHS)} + + fig, axes = plt.subplots(3, 4, figsize=(15, 9.5), sharex=True, sharey=True) + for r, e in enumerate(exps): + for c, tag in enumerate(RH_TAGS): + pm = banana_panel(axes[r, c], runs[(e, tag)], t_hr, + first=(r, c) == (0, 0)) + pm.set_norm(LogNorm(vmin=1e0, vmax=1e5)) + axes[r, c].set_xscale("log") + axes[r, c].set_xlim(t_hr[0], 12.0) + if r == 0: + axes[r, c].set_title(RH_LBL[tag], loc="left") + if c == 0: + axes[r, c].set_ylabel(exp_lbl[e] + "\ndry $D_p$ (nm)", fontsize=8.5) + if r == 2: + axes[r, c].set_xlabel("time (h, log)") + fig.colorbar(pm, ax=axes, label=r"dN/dln$D_p$ (# mg$^{-1}$)", pad=0.015) + fig.legend(*axes[0, 0].get_legend_handles_labels(), loc="lower center", + ncol=5, frameon=False, bbox_to_anchor=(0.45, -0.03), + title="mode $d_{gn}$ overlays") + fig.suptitle("Banana plots (log time) — stratospheric RH sweep, 12 h", y=0.995) + fig.savefig(out / "cam_strat_rh_banana_12h.png", dpi=150, bbox_inches="tight") + plt.close(fig) + + num0, dgn0 = init_dist_inputs() + snap_h = [0.0, 1.0, 3.0, 6.0, 12.0] + fig, axes = plt.subplots(3, 4, figsize=(15, 9.5), sharex=True, sharey=True) + for r, e in enumerate(exps): + for c, tag in enumerate(RH_TAGS): + ax = axes[r, c] + rd = runs[(e, tag)] + for h, col in zip(snap_h, GREYS): + if h == 0.0: + tot, _ = dist(num0, dgn0) + lbl = "initial" + else: + it = int(h * 3600 / DT) - 1 + tot, _ = dist(rd["num"][it], rd["dgn"][it]) + lbl = f"{h:g} h" + ax.plot(DP * 1e9, tot, color=col, lw=1.5, + label=lbl if (r, c) == (0, 0) else None) + totw, _ = dist(rd["num"][-1], rd["dgnwet"][-1]) + ax.plot(DP * 1e9, totw, color="#1c5cab", lw=1.5, ls="--", + label="12 h, WET $D_p$" if (r, c) == (0, 0) else None) + ax.set_xscale("log"); ax.set_yscale("log") + ax.set_xlim(1, 3000); ax.set_ylim(1e-2, 3e5) + ax.grid(True, which="major", **GRID) + if r == 0: + ax.set_title(RH_LBL[tag], loc="left") + if c == 0: + ax.set_ylabel(exp_lbl[e] + "\n" + r"dN/dln$D_p$ (# mg$^{-1}$)", + fontsize=8.5) + if r == 2: + ax.set_xlabel("$D_p$ (nm)") + fig.legend(*axes[0, 0].get_legend_handles_labels(), loc="lower center", + ncol=6, frameon=False, bbox_to_anchor=(0.5, -0.025)) + fig.suptitle("Snapshots (solid grey = dry, by time; dashed blue = wet at 12 h)", + y=0.995) + fig.tight_layout() + fig.savefig(out / "cam_strat_rh_snapshots_12h.png", dpi=150, bbox_inches="tight") + plt.close(fig) + + fig, axes = plt.subplots(3, 4, figsize=(15, 9.0)) + for r, e in enumerate(exps): + axN, axM, axG, axD = axes[r] + for tag, col in zip(RH_TAGS, RH_C): + rd = runs[(e, tag)] + axN.plot(t_hr, rd["num"].sum(1) * 1e-6, color=col, lw=1.7, + label=RH_LBL[tag]) + axM.plot(t_hr, rd["so4"].sum(1) * 1e9, color=col, lw=1.7) + axG.plot(t_hr, rd["h2so4"], color=col, lw=1.7) + axD.plot(t_hr, rd["dgnwet"][:, 1] * 1e9, color=col, lw=1.7) + axD.plot(t_hr, runs[(e, "rh50")]["dgn"][:, 1] * 1e9, color="#888888", + lw=1.3, ls=":", label="dry (RH 50%)" if r == 0 else None) + axN.set_yscale("log"); axG.set_yscale("log") + axN.set_ylabel(exp_lbl[e] + "\n" + r"total N (# mg$^{-1}$)", fontsize=8.5) + axM.set_ylabel(r"total SO$_4$ (µg kg$^{-1}$)") + axG.set_ylabel(r"H$_2$SO$_4$(g) (molec cm$^{-3}$)") + axD.set_ylabel("aitken $D_{gn}$ (nm)") + for ax in axes[r]: + ax.grid(True, **GRID) + if r == 2: + ax.set_xlabel("time (h)") + axes[0, 0].set_title("(a) total number", loc="left") + axes[0, 1].set_title("(b) total sulfate aerosol mass", loc="left") + axes[0, 2].set_title("(c) H$_2$SO$_4$ gas", loc="left") + axes[0, 3].set_title("(d) aitken wet diameter (dotted grey = dry)", loc="left") + axes[0, 0].legend(fontsize=8, framealpha=0.95) + axes[0, 3].legend(fontsize=8, framealpha=0.95) + fig.suptitle("Budgets — stratospheric RH sweep, 12 h", y=0.998) + fig.tight_layout() + fig.savefig(out / "cam_strat_rh_budgets_12h.png", dpi=150, bbox_inches="tight") + plt.close(fig) + + +if __name__ == "__main__": + out = Path(sys.argv[1]) if len(sys.argv) > 1 else ( + Path(__file__).resolve().parents[1] / "docs" / "figures") + out.mkdir(parents=True, exist_ok=True) + _style() + fig_sweep_6h(out) + print("6-h sweep figures written") + fig_experiments_12h(out) + print("12-h experiment figures written") diff --git a/tests/test_cam_driver.py b/tests/test_cam_driver.py index 1b94b8f..9e0148d 100644 --- a/tests/test_cam_driver.py +++ b/tests/test_cam_driver.py @@ -168,3 +168,44 @@ def test_documented_defaults() -> None: sig = inspect.signature(cd.cam_run_step) assert sig.parameters["n_substeps"].default == 16 assert sig.parameters["reseed_dgnwet_each_step"].default is True + + +# --------------------------------------------------------------------------- +# plan 027 — jit/scan semantics +# --------------------------------------------------------------------------- + +def test_scan_trajectory_is_bit_identical_to_stepping() -> None: + """cam_run_timesteps' scan inlines the SAME jitted step body that + cam_run_step dispatches, so the trajectory must equal per-step + calls exactly — not just approximately.""" + topo = CAM_MAM4 + s = build_ic(topo) + for i in range(3): + s = cd.cam_run_step(s, topology=topo, strat=True, n_substeps=4, + first_step=(i == 0)) + _, traj = cd.cam_run_timesteps(build_ic(topo), 3, topology=topo, + strat=True, n_substeps=4) + np.testing.assert_array_equal(np.asarray(traj["q"][-1]), + np.asarray(s["q"])) + + +def test_rate_is_traced_and_substeps_are_static() -> None: + """ADR-020's split, enforced: a different so2_to_h2so4_rate VALUE + must reuse the compilation (traced leaf — a rate sweep must not + recompile), while a different n_substeps must get its own cache + entry (it changes the scan length — genuinely different code).""" + topo = CAM_MAM4 + ic = build_ic(topo) + cd._cam_run_step_jit.clear_cache() + out_a = cd.cam_run_step(ic, topology=topo, n_substeps=2, + so2_to_h2so4_rate=1.0e-5) + size_1 = cd._cam_run_step_jit._cache_size() + out_b = cd.cam_run_step(ic, topology=topo, n_substeps=2, + so2_to_h2so4_rate=2.0e-5) + assert cd._cam_run_step_jit._cache_size() == size_1, ( + "a new rate VALUE triggered a recompile — it must be traced") + assert not np.allclose(np.asarray(out_a["q"]), np.asarray(out_b["q"])) + cd.cam_run_step(ic, topology=topo, n_substeps=3, + so2_to_h2so4_rate=1.0e-5) + assert cd._cam_run_step_jit._cache_size() == size_1 + 1, ( + "a new n_substeps did not retrace — it must be static")