Skip to content

Make the pointwise 2M+P3 tendencies ForwardDiff-differentiable#725

Open
haakon-e wants to merge 2 commits into
mainfrom
he/ad-compat-2m-p3
Open

Make the pointwise 2M+P3 tendencies ForwardDiff-differentiable#725
haakon-e wants to merge 2 commits into
mainfrom
he/ad-compat-2m-p3

Conversation

@haakon-e

@haakon-e haakon-e commented Jun 11, 2026

Copy link
Copy Markdown
Member

Implicit (Rosenbrock-type) microphysics substepping needs Jacobians of the pointwise 2M+P3 tendencies w.r.t. the 8 prognostic species, so ForwardDiff.jacobian must work through bulk_microphysics_tendencies (logλ is held fixed). This PR does some groundwork to enable this.

  • Modify P3State and Γ_incl to accept a mix of Dual and floating point arguments.
  • Implement Dual method for SF.gamma_inc_inv (What about gamma_inc_inv's derivative ? JuliaMath/SpecialFunctions.jl#452) locally. Ideally this should be added to SpecialFunctions, and I might do so later, but adding here to move forward locally for now. The implementation uses a mix of analytical and central differences of the forward map.
  • Ensure ice_mass_coeffs returns one tuple type (needed for efficient AD).
  • Type early returns and fallback values across the 2M, NonEq, and P3 leaves by the promotion of all the arguments the result derives from (new Utilities.promote_typeof), so a caller mixing Dual and float arguments does not get a union-typed, heap-boxed return.
  • Handle edge cases for sgs_weight_function (lower tail rounds 1 - a to 1, so atanh(-1) = -Inf) and SB2006 autoconversion (vertical tangent of τ^0.7 at zero rain), which previously could produce NaNs.
  • Unrelated fix: get_distribution_logλ_from_prognostic called an undefined function.

I might try to differentiate through the logλ shape solve later; which needs ∂/∂a of the forward gamma_inc (JuliaMath/SpecialFunctions.jl#317).

Tests cover Jacobians at four regimes for both float types against per-row central finite differences, the gamma_inc_inv rule into the q = eps() tail, the regularised-ratio band, and that the touched functions stay concretely typed under single-Dual and all-Dual argument mixes.

`get_state_from_prognostic` is not defined anywhere; the path was
unexercised by tests.
@haakon-e haakon-e force-pushed the he/ad-compat-2m-p3 branch 3 times, most recently from 4759ee8 to 374c47c Compare June 11, 2026 20:14
@codecov

codecov Bot commented Jun 11, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.92%. Comparing base (ea430c0) to head (aa342a2).

Additional details and impacted files
@@            Coverage Diff             @@
##             main     #725      +/-   ##
==========================================
+ Coverage   92.74%   92.92%   +0.18%     
==========================================
  Files          56       56              
  Lines        2701     2728      +27     
==========================================
+ Hits         2505     2535      +30     
+ Misses        196      193       -3     
Components Coverage Δ
src 93.77% <100.00%> (+0.17%) ⬆️
ext 69.47% <ø> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Implicit substepping needs Jacobians of the pointwise tendencies
w.r.t. the 8 prognostic species (logλ held fixed) — mixed Dual/plain
arguments, frozen thermodynamic context. Primal values and
allocations are unchanged (uniform-float promotion is the identity).

- Type flexibility for mixed Dual/plain inputs: P3State, Γ_incl, the
  gamma_inc_inv wrapper, ice_mass_coeffs (uniform branch return —
  was boxing ~105 MiB per Jacobian), logistic_function_integral,
  Chen2022_exponential_pdf, LclRaiRates, the size-distribution
  closures.
- gamma_inc_inv: add an implicit-function-theorem Dual rule (none
  upstream), exact in the quantile slots; the shape slot differences
  the smaller of (P, Q), the larger saturating in the tails.
- SB2006 autoconversion: ϕ_au ∝ τ^0.7 has a vertical tangent at
  τ = 0, so ∂/∂q_rai is NaN at zero rain with cloud present. Gate
  the enhancement below ϵ_numerics_2M_M; value at q_rai = 0 unchanged.
- sgs_weight_function: guard the lower tail where atanh → -Inf
  poisons Float32 rime-ratio partials.
- Sentinel returns: typing from one argument (eltype(q_tot)-style)
  makes the return union-typed and heap-boxed under mixed-type
  callers — silently, since the box lives in the abstract return
  convention (JET and the optimized IR show nothing). Add
  Utilities.promote_typeof and apply it across the verified hits in
  the 2M, NonEq, P3, and BMT code.

Tests: per-row Jacobian-vs-FD oracle; AD coverage for the
regularised-ratio band, the tail shape slots, and mixed arguments;
sentinel_type_tests assert concrete returns per signature.

Deferred: the logλ shape solve (needs ∂/∂a of forward gamma_inc) and
the 1M sentinel leaves.
@haakon-e haakon-e force-pushed the he/ad-compat-2m-p3 branch from 415fe44 to b150213 Compare June 14, 2026 02:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant