Make the pointwise 2M+P3 tendencies ForwardDiff-differentiable by haakon-e · Pull Request #725 · CliMA/CloudMicrophysics.jl

haakon-e · 2026-06-11T19:57:31Z

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.

codecov · 2026-06-11T20:29:18Z

Codecov Report

❌ Patch coverage is 99.06542% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 92.96%. Comparing base (f68f56f) to head (ddd64f0).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #725      +/-   ##
==========================================
+ Coverage   92.68%   92.96%   +0.28%     
==========================================
  Files          56       56              
  Lines        2816     2859      +43     
==========================================
+ Hits         2610     2658      +48     
+ Misses        206      201       -5

Components	Coverage Δ
src	`93.77% <99.06%> (+0.28%)`	⬆️
ext	`69.47% <ø> (ø)`

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

haakon-e · 2026-06-12T15:31:53Z

This change is part of the following stack:

Make the pointwise 2M+P3 tendencies ForwardDiff-differentiable #725 ◀
- feat(P3): Gauss-Legendre quadrature, integrand breakpoints -> GL(6) default #741
- feat(Rosenbrock): unified RosenbrockAverage substepping framework + 2M+P3 exact-AD mode #726
  - feat(Rosenbrock): 1M donor / coupled / exact modes on the unified framework #731
    - feat(Rosenbrock): post-solve per-process verbose diagnostics #732
      - fix(Rosenbrock): saturation adjustment on max(S_ice, S_liq) #737
        
        perf(BMT): recover the primal tendency from the Rosenbrock Jacobian pass #742
        
        feat(BMT): analytic ManualJacobian for the 2M+P3 Rosenbrock substep #743
        
        feat(BMT): coupled Tier-3 donor linearization for the 2M+P3 ManualJacobian #744

_{Change managed by git-spice.}

`get_state_from_prognostic` is not defined anywhere; the path was unexercised by tests.

Implicit substepping needs Jacobians of the pointwise tendencies w.r.t. the 8 prognostic species (`logλ` held fixed). - Type flexibility for mixed Dual/plain inputs: P3State, Γ_incl, the gamma_inc_inv wrapper, ice_mass_coeffs, 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. Condition the enhancement below ϵ_numerics_2M_M; value at q_rai = 0 unchanged. - sgs_weight_function: Check the lower tail where atanh → -Inf makes rime-ratio partials bad in Float32. - Early returns: typing from one argument (e.g. `FT = eltype(q_tot)`) makes the return union-typed and heap-boxed for mixed-type callers. Add Utilities.promote_typeof and apply it where relevant. Tests: per-row Jacobian-vs-FD consistency tests; 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 early return cases.

The analytical solution for ρ_d has a denominator that is 0/0 as F_rim → 0. In Float32 the cancellation drives ρ_d, and therefore ρ_g, negative, which later raises a DomainError in get_D_gr. Rewrite get_ρ_d with the relative exponential functions exprel1 and exprel2 so the cancellation is removed; ρ_g then stays positive for every physical input, with no clipping. Add the derivation and an error-cancellation figure to the P3 documentation, and a Float32 stability test.

The differentiated 2M+P3 tendency must compile for CUDA. Several constructs that are fine on the CPU break the GPU compile by heap-allocating or boxing: - get_ρ_d / exponential moments: avoid the host factorial table and Rational exponents (use a `fac` helper and float coefficients). - P3State construction: convert the derived fields with explicit `FT(...)` calls instead of broadcasting the (parametric, possibly `Dual`) element type over a tuple, which the GPU back end heap-allocates. - liquid_freezing_rate: type the fallback zero via `promote_typeof` so the `ifelse` branches share an element type under ForwardDiff. - gamma_inc_inv: take the inverse-CDF shape derivative with a one-sided finite difference at a fixed step, removing one `gamma_inc` evaluation.

trontrytel · 2026-07-01T22:36:05Z

 """
 function liquid_freezing_rate(opt::CMP.RainFreezing, pdf, tps, q, ρ, N, T)
-    FT = eltype(q)
+    FT = float(UT.promote_typeof(q, ρ, N, T))


How expensive it is per element in the list here? Does it save us any to just have q as before?

I'm not quite sure what you mean by expensive? I think the "proper" fix instead of this is to merge #746, as we discussed in the office, where FT is derived from the "return type" of the relevant calculation.

Though that PR is currently implemented as cond ? f() : otherwise(FT), so if we're supposed to rewrite everything with ifelse as you mentioned elsewhere, I'd have to update the PR first.

I meant float(UT.promote_typeof(q, ρ, N, T)) vs float(UT.promote_typeof(q) Since I expect it's q is the one with the potentially tricky type.

And yeah, I wasn't sure how orthogonal our discussion is compared to the GPU performance PR

trontrytel · 2026-07-01T22:37:38Z

    # Avoid NaN when both q and x_min are 0
-    N_max = iszero(x_min) ? oftype(q, Inf) : ρ * q / x_min
+    FT = UT.promote_typeof(q, ρ, N)
+    N_max = iszero(x_min) ? FT(Inf) : FT(ρ * q / x_min)


Are we supposed to use ifelse now?

I suspect the compiler will internally rewrite these to ifelse, but can do it explicitly here for extra safety:)

trontrytel · 2026-07-01T22:39:24Z

+"""
+    fac(n)
+
+Integer factorial `n!`, valid for `0 ≤ n ≤ 20`.


Should we throw an error for n>20 ?

sure, I can do that. In practice, we'll never use this for numbers even close to that anyways.

So we think

trontrytel · 2026-07-01T22:47:19Z

+        # fallback values typed by the promotion of the node and the captured state
+        # (mixed plain/Dual under differentiation)
+        FT = UT.promote_typeof(Dᵢ, ΔT, Δρᵥ_sat, denom)
+        Tₐ ≥ T_frz && return zero(FT)  # No collisional freezing above the freezing temperature


Are we supposed to do ifelse now?

thanks for the catch! I'll update these.

haakon-e · 2026-07-01T22:49:05Z

welp. Sorry @trontrytel I'll put your feedback into the next PR. I forgot I had auto merge on!

trontrytel · 2026-07-01T22:52:07Z

welp. Sorry @trontrytel I'll put your feedback into the next PR. I forgot I had auto merge on!

No worries! I knowingly hit approve

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

haakon-e force-pushed the he/ad-compat-2m-p3 branch from 374c47c to bdd359d Compare June 11, 2026 20:53

This was referenced Jun 12, 2026

feat(Rosenbrock): unified RosenbrockAverage substepping framework + 2M+P3 exact-AD mode #726

Open

Type 2M sentinel returns by all result-feeding inputs #727

Closed

haakon-e force-pushed the he/ad-compat-2m-p3 branch 2 times, most recently from 7514748 to 415fe44 Compare June 14, 2026 02:02

haakon-e force-pushed the he/ad-compat-2m-p3 branch from 415fe44 to b150213 Compare June 14, 2026 02:45

haakon-e mentioned this pull request Jun 17, 2026

fix(Rosenbrock): saturation adjustment on max(S_ice, S_liq) #737

Draft

haakon-e force-pushed the he/ad-compat-2m-p3 branch from 643b5a8 to aef2c68 Compare June 18, 2026 22:41

haakon-e changed the base branch from main to he/rft-cleanup-from-681 June 18, 2026 22:42

haakon-e mentioned this pull request Jun 18, 2026

rft: cleanup from #681 #736

Merged

haakon-e force-pushed the he/rft-cleanup-from-681 branch from 31e584b to ea7a9b0 Compare June 22, 2026 21:00

haakon-e force-pushed the he/ad-compat-2m-p3 branch from aef2c68 to 4696205 Compare June 22, 2026 21:01

haakon-e force-pushed the he/rft-cleanup-from-681 branch 2 times, most recently from 55f2d59 to 09d92b0 Compare June 22, 2026 21:53

haakon-e force-pushed the he/ad-compat-2m-p3 branch from 4696205 to 687d8e7 Compare June 22, 2026 23:42

Base automatically changed from he/rft-cleanup-from-681 to main June 22, 2026 23:49

haakon-e force-pushed the he/ad-compat-2m-p3 branch from 687d8e7 to db7adcf Compare June 22, 2026 23:55

This was referenced Jun 23, 2026

feat(P3): Gauss-Legendre quadrature, integrand breakpoints -> GL(6) default #741

Open

fix(P3): use per-regime material density in the ice aspect-ratio closure #740

Closed

haakon-e force-pushed the he/ad-compat-2m-p3 branch from 26145ad to 2c7d843 Compare June 23, 2026 22:26

trontrytel reviewed Jun 24, 2026

View reviewed changes

Comment thread src/P3_size_distribution.jl Outdated

haakon-e mentioned this pull request Jun 24, 2026

feat(Utilities): add safe_call for type-inferred conditional fallbacks #746

Closed

haakon-e marked this pull request as draft June 25, 2026 18:00

haakon-e force-pushed the he/ad-compat-2m-p3 branch from 6436c12 to 710d1ce Compare June 25, 2026 18:01

haakon-e marked this pull request as ready for review June 25, 2026 18:01

haakon-e marked this pull request as draft June 25, 2026 18:19

haakon-e marked this pull request as ready for review June 26, 2026 00:00

haakon-e marked this pull request as draft June 26, 2026 00:00

haakon-e force-pushed the he/ad-compat-2m-p3 branch 3 times, most recently from 98e82b8 to 6878dd5 Compare June 26, 2026 01:31

haakon-e marked this pull request as ready for review June 26, 2026 18:32

haakon-e added 3 commits June 30, 2026 22:14

Fix undefined call in get_distribution_logλ_from_prognostic

0de32d9

`get_state_from_prognostic` is not defined anywhere; the path was unexercised by tests.

haakon-e force-pushed the he/ad-compat-2m-p3 branch 2 times, most recently from 044509d to 11eb605 Compare July 1, 2026 05:45

haakon-e force-pushed the he/ad-compat-2m-p3 branch from 11eb605 to ddd64f0 Compare July 1, 2026 06:24

haakon-e enabled auto-merge July 1, 2026 16:57

trontrytel reviewed Jul 1, 2026

View reviewed changes

trontrytel approved these changes Jul 1, 2026

View reviewed changes

haakon-e merged commit 5b4b2f5 into main Jul 1, 2026
17 of 20 checks passed

haakon-e deleted the he/ad-compat-2m-p3 branch July 1, 2026 22:48

Uh oh!

Conversation

haakon-e commented Jun 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Jun 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

haakon-e commented Jun 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

haakon-e commented Jul 1, 2026

Uh oh!

trontrytel commented Jul 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

haakon-e commented Jun 11, 2026 •

edited

Loading

codecov Bot commented Jun 11, 2026 •

edited

Loading

haakon-e commented Jun 12, 2026 •

edited

Loading