Wip on Jacobian derivatives

[trontrytel]

This PR adds tendency derivatives for the 1-moment microphysics and non-equilibrium cloud formation. They are to be used in the implicit solver in the Atmos model.

• For cloud condensation/evaporation and deposition/sublimation we provide a full derivative and a simplified one equal to 1/timescale
• We neglect the derivatives for collision terms like autoconversion and accretion
• We model the snow melt, snow deposition/sublimation and rain evaporation derivatives as source/q

For each of the above processes the function now returns the rate and rate derivative. The Bulk microphysics tendency also returns the sum of tendencies and the sum of derivatives.

I have added to the documentation pages and included some plots. I have also added to the tests, including some tests based on finite-difference "manual" calculation of the derivatives.

Going forward, it would not be hard to add more analytical terms for the derivatives. Or, in case we want to make the derivative free of having to redo the math if we change the parameterization functions, I could code all of those as a finite difference numerical computation (something like deriv = (micro_func(q+eps) - micro_func(q)) / eps)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.80%. Comparing base (ed3c40a) to head (d7f354a).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #682      +/-   ##
==========================================
+ Coverage   94.74%   94.80%   +0.05%     
==========================================
  Files          53       53              
  Lines        2246     2270      +24     
==========================================
+ Hits         2128     2152      +24     
  Misses        118      118              

Components	Coverage Δ
src	95.90% <100.00%> (+0.04%)	⬆️
ext	69.47% <ø> (ø)

🚀 New features to boost your workflow:

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

[trontrytel]

I still need to check some signs. But please lmk if this is what we need.

[tapios]

Looks like a great start! In the tests, we need to be careful to evaluate the derivatives at fixed e_tot and q_tot (and also make clear in the docs that this is what it is meant to be).

[tapios]

You could use TD function for saturation specific humidity directly.

[tapios]

Do tighter bounds work?

[tapios]

Same here: Do tighter bounds work?

[tapios]

This isn't testing quite the right derivative. We need to get the derivative at fixed q_tot (which you do here) and e_tot (which is not done here, instead T is fixed).

For the correct test, we'd need to wrap this in a function that keeps e_tot fixed but updates T.

[tapios]

Same as above: we need to test at fixed e_tot, not fixed temperature.

[tapios]

Same as above.

[tapios]

That's a better, taking first-order temperature change into account. Would be better to hold e_tot exactly fixed in a wrapper.

[haakon-e]

Thanks for putting this PR together! Writing down these derivatives is a good step toward a reasonable implicit solver.

I do have some concerns with this design, in which all the process tendencies return a tuple:

• With the method names as is, I don't think it's intuitive that they now also return derivatives
• Different methods return different-length tuples
• it's not necessarily intuitive with respect to which variable(s) the derivatives are taken (unless you inspect the source code / docstring)

I would much prefer keeping the tendency methods as-is, and define new methods for the derivatives. For example:

# src/Microphysics1M.jl
function ∂accretion_∂q_lcl(::CMP.Acnv1M{FT}), q_lcl, smooth_transition = false) where {FT}
    return FT(0)
end
# ...

function ∂evaporation_sublimation_∂q_rai(
    ::CMP.Rain{FT}, 
    ::CMP.Blk1MVelTypeRain{FT},
    q_rai,
    tendency,
) where {FT}
    # derivative approximation: ∂(tendency)/∂q_rai ≈ tendency / q_rai
    return q_rai > UT.ϵ_numerics(FT) ? tendency / q_rai : zero(tendency)
end

# src/MicrophysicsNonEq.jl
∂conv_q_vap_to_q_lcl_icl_MM2015_∂q_cond_simple(p::CMP.CloudCondensateType) = -1 / τ_relax(p)

saturation_vapor_pressure(::CMP.CloudLiquid, tps::TDI.PS, T) =
    TDI.saturation_vapor_pressure_over_liquid(tps, T)
saturation_vapor_pressure(::CMP.CloudIce, tps::TDI.PS, T) =
    TDI.saturation_vapor_pressure_over_ice(tps, T)

saturation_specific_humidity(p::CMP.CloudCondensateType, tps::TDI.PS, ρ, T) =
    TDI.p2q(tps, T, ρ, saturation_vapor_pressure(p, tps, T))

function ∂²qᵥ_sat_∂T²(tps::TDI.PS, p::CMP.CloudCondensateType, ρ, T)
    Rᵥ = TDI.Rᵥ(tps)
    Lᵥ = TDI.Lᵥ(tps, T)
    qᵥ_sat = saturation_specific_humidity(p, tps, ρ, T)
    return qᵥ_sat * ((Lᵥ / Rᵥ / T^2 - 1/T)^2 + (1/T^2 - 2 * Lᵥ / Rᵥ / T^3))
end

function ∂qᵥ_sat_∂T(tps::TDI.PS, p::CMP.CloudCondensateType, ρ, T)
    Rᵥ = TDI.Rᵥ(tps)
    Lᵥ = TDI.Lᵥ(tps, T)
    qᵥ_sat = saturation_specific_humidity(p, tps, ρ, T)
    return qᵥ_sat * (Lᵥ / (Rᵥ * T^2) - 1 / T)
end

function ∂conv_q_vap_to_q_lcl_icl_MM2015_∂q_cond(
    p::CMP.CloudCondensateType, tps::TDI.PS,
    q_tot, q_lcl, q_icl, q_rai, q_sno, ρ, T,
    tendency,
)
    ∂qᵥs_dT = ∂qᵥ_sat_∂T(tps, p, ρ, T)
    ∂²qᵥs_∂T² = ∂²qᵥ_sat_∂T²(tps, p, ρ, T)
    cₚ_air = TDI.cpₘ(tps, q_tot, q_lcl + q_rai, q_icl + q_sno)
    Γ = 1 + (Lᵥ / cₚ_air) * ∂qᵥs_dT
    return - 1 / τ_relax(p) - tendency / Γ * (Lᵥ / cₚ_air)^2 * ∂²qᵥs_∂T²
end

[tapios]

Thanks for putting this PR together! Writing down these derivatives is a good step toward a reasonable implicit solver.

I do have some concerns with this design, in which all the process tendencies return a tuple:

• With the method names as is, I don't think it's intuitive that they now also return derivatives
• Different methods return different-length tuples
• it's not necessarily intuitive with respect to which variable(s) the derivatives are taken (unless you inspect the source code / docstring)

I would much prefer keeping the tendency methods as-is, and define new methods for the derivatives. For example:

# src/Microphysics1M.jl
function ∂accretion_∂q_lcl(::CMP.Acnv1M{FT}), q_lcl, smooth_transition = false) where {FT}
    return FT(0)
end
# ...

function ∂evaporation_sublimation_∂q_rai(
    ::CMP.Rain{FT}, 
    ::CMP.Blk1MVelTypeRain{FT},
    q_rai,
    tendency,
) where {FT}
    # derivative approximation: ∂(tendency)/∂q_rai ≈ tendency / q_rai
    return q_rai > UT.ϵ_numerics(FT) ? tendency / q_rai : zero(tendency)
end

# src/MicrophysicsNonEq.jl
∂conv_q_vap_to_q_lcl_icl_MM2015_∂q_cond_simple(p::CMP.CloudCondensateType) = -1 / τ_relax(p)

saturation_vapor_pressure(::CMP.CloudLiquid, tps::TDI.PS, T) =
    TDI.saturation_vapor_pressure_over_liquid(tps, T)
saturation_vapor_pressure(::CMP.CloudIce, tps::TDI.PS, T) =
    TDI.saturation_vapor_pressure_over_ice(tps, T)

saturation_specific_humidity(p::CMP.CloudCondensateType, tps::TDI.PS, ρ, T) =
    TDI.p2q(tps, T, ρ, saturation_vapor_pressure(p, tps, T))

function ∂²qᵥ_sat_∂T²(tps::TDI.PS, p::CMP.CloudCondensateType, ρ, T)
    Rᵥ = TDI.Rᵥ(tps)
    Lᵥ = TDI.Lᵥ(tps, T)
    qᵥ_sat = saturation_specific_humidity(p, tps, ρ, T)
    return qᵥ_sat * ((Lᵥ / Rᵥ / T^2 - 1/T)^2 + (1/T^2 - 2 * Lᵥ / Rᵥ / T^3))
end

function ∂qᵥ_sat_∂T(tps::TDI.PS, p::CMP.CloudCondensateType, ρ, T)
    Rᵥ = TDI.Rᵥ(tps)
    Lᵥ = TDI.Lᵥ(tps, T)
    qᵥ_sat = saturation_specific_humidity(p, tps, ρ, T)
    return qᵥ_sat * (Lᵥ / (Rᵥ * T^2) - 1 / T)
end

function ∂conv_q_vap_to_q_lcl_icl_MM2015_∂q_cond(
    p::CMP.CloudCondensateType, tps::TDI.PS,
    q_tot, q_lcl, q_icl, q_rai, q_sno, ρ, T,
    tendency,
)
    ∂qᵥs_dT = ∂qᵥ_sat_∂T(tps, p, ρ, T)
    ∂²qᵥs_∂T² = ∂²qᵥ_sat_∂T²(tps, p, ρ, T)
    cₚ_air = TDI.cpₘ(tps, q_tot, q_lcl + q_rai, q_icl + q_sno)
    Γ = 1 + (Lᵥ / cₚ_air) * ∂qᵥs_dT
    return - 1 / τ_relax(p) - tendency / Γ * (Lᵥ / cₚ_air)^2 * ∂²qᵥs_∂T²
end

Along the same lines: We will need the derivatives at different call sites (Jacobian) than the tendencies themselves. So it's probably better, and better to avoid duplicate computations, to have separate functions for the derivatives.

Externally, we will only need the fused tendencies and their approximate derivatives, which it would be good to have as separate functions.