diff --git a/Project.toml b/Project.toml index 647a08f0..a9f0aa97 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Thermodynamics" uuid = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c" authors = ["Climate Modeling Alliance"] -version = "0.15.7" +version = "0.15.8" [deps] ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" diff --git a/docs/src/API.md b/docs/src/API.md index e9181ae2..56645cb9 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -125,6 +125,12 @@ q_vap_saturation_from_pressure ```@docs saturation_adjustment +∂e_int_∂T_sat_ρ +∂e_int_∂T_sat_p +∂θ_li_∂T_sat_ρ +∂θ_li_∂T_sat_p +∂p_∂T_sat_ρ +∂h_∂T_sat_p ``` ### Air Entropies diff --git a/docs/src/HowToGuide.md b/docs/src/HowToGuide.md index 9e1ff14c..4e27d902 100644 --- a/docs/src/HowToGuide.md +++ b/docs/src/HowToGuide.md @@ -87,7 +87,6 @@ sol = TD.saturation_adjustment( T_equil = sol.T q_liq_equil = sol.q_liq q_ice_equil = sol.q_ice -converged = sol.converged println("Equilibrium T: $T_equil") ``` diff --git a/src/saturation_adjustment.jl b/src/saturation_adjustment.jl index 2f720b12..cf6dfdba 100644 --- a/src/saturation_adjustment.jl +++ b/src/saturation_adjustment.jl @@ -3,6 +3,16 @@ import RootSolvers as RS export saturation_adjustment +export ∂e_int_∂T_sat_ρ +export ∂e_int_∂T_sat_p +export ∂h_∂T_sat_p +export ∂θ_li_∂T_sat_ρ +export ∂θ_li_∂T_sat_p +export ∂p_∂T_sat_ρ + +# --------------------------------------------- +# Public API: full solver methods +# --------------------------------------------- """ saturation_adjustment( @@ -45,8 +55,7 @@ given density `ρ`, internal energy `e_int`, and total specific humidity `q_tot` # Notes - This function solves for `T` such that `e_int = internal_energy_sat(param_set, T, ρ, q_tot)` using root-finding, then computes `(q_liq, q_ice)` from [`condensate_partition`](@ref). -- For `ρe` formulation, `NewtonsMethod` is recommended (fast + analytic derivative available). -- For other formulations, `SecantMethod` or `BrentsMethod` are recommended. +- `NewtonsMethod` is recommended for all formulations (analytic derivatives available). - **GPU broadcasting**: Pass `forced_fixed_iters` as a positional Bool. """ function saturation_adjustment( @@ -62,13 +71,7 @@ function saturation_adjustment( forced_fixed_iters::Bool = false, ) where {M} if forced_fixed_iters - return saturation_adjustment_ρe_fixed_iters( - param_set, - ρ, - e_int, - q_tot, - maxiter, - ) + return saturation_adjustment_fixed_iters(param_set, ρe(), ρ, e_int, q_tot, maxiter) end temp_unsat_func = (param_set, e_int, q_tot) -> air_temperature(param_set, e_int, q_tot) @@ -80,7 +83,7 @@ function saturation_adjustment( air_temperature(param_set, ρe(), e_int, q_tot, zero(q_tot), q_tot) # Root function (supports analytic derivative if M=NewtonsMethod) - roots_func = _make_roots_function(M, param_set, ρ, e_int, q_tot) + roots_func = _make_roots_function(M, param_set, ρe(), ρ, e_int, q_tot) (T, converged) = _saturation_adjustment_generic( M, @@ -100,149 +103,6 @@ function saturation_adjustment( return (; T, q_liq, q_ice, converged) end -""" - saturation_adjustment_ρe_fixed_iters(param_set, ρ, e_int, q_tot, maxiter) - -GPU-optimized saturation adjustment for `ρe` formulation using a fixed number of -Newton iterations. - -Bypasses standard solver logic (bracketing, unsaturated checks, convergence testing) -to avoid branch divergence on GPUs. - -# Arguments -- `param_set`: Thermodynamics parameter set -- `ρ`: Density [kg/m³] -- `e_int`: Specific internal energy [J/kg] -- `q_tot`: Total specific humidity [kg/kg] -- `maxiter`: Number of Newton iterations (2 recommended for < 0.1 K accuracy) - -# Returns -- `NamedTuple` `(; T, q_liq, q_ice, converged)` - -# Notes -- The `converged` field is always `true` (no convergence check is performed). -- With `maxiter = 2`, temperature accuracy is better than 0.1 K for typical atmospheric - conditions (T < 320 K). -- This is an internal helper function. For the public API, use [`saturation_adjustment`](@ref) - with `forced_fixed_iters=true` as a positional argument. -""" -@inline function saturation_adjustment_ρe_fixed_iters( - param_set::APS, - ρ, - e_int, - q_tot, - maxiter, -) - # Estimate T_unsat for initialization. - T_unsat = air_temperature(param_set, e_int, q_tot) - T_init_min = TP.T_init_min(param_set) - T = max(T_init_min, T_unsat) - - @fastmath for _ in 1:maxiter - e_val = internal_energy_sat(param_set, T, ρ, q_tot) - de_int_dT = ∂e_int_∂T_sat(param_set, T, ρ, q_tot) - - # f(T) = e_int_target - e_val - # f'(T) = -de_int_dT - # Newton step: T_{n+1} = T_n - f(T_n) / f'(T_n) - # delta = - (e_int_target - e_val) / (-de_int_dT) = (e_int_target - e_val) / de_int_dT - T += (e_int - e_val) / de_int_dT - end - - (q_liq, q_ice) = condensate_partition(param_set, T, ρ, q_tot) - return (; T, q_liq, q_ice, converged = true) -end -""" - bound_upper_temperature(param_set, T_lo, T_hi) - -Internal function. Bounds the upper temperature guess `T_hi` for bracket methods. - -Returns `T_hi` bounded by `T_max`, ensuring it is at least `T_lo + 0.1` for -valid numerical initialization. -""" -@inline function bound_upper_temperature(param_set, T_lo, T_hi) - FT = eltype(param_set) - T_max = TP.T_max(param_set) - # Ensure T_hi is physically valid (<= T_max) - T_hi_phys = min(T_max, T_hi) - # Ensure T_hi > T_lo for numerical initialization (use relative tolerance) - return max(T_lo * (1 + FT(1e-3)), T_hi_phys) -end - - -""" - _make_sa_solver(::Type{M}, param_set, T_unsat, T_ice, T_guess) - -Internal helper to construct a root-solving method instance for saturation adjustment. - -# Arguments -- `M`: Root-solving method type (e.g., `RS.NewtonsMethod`, `RS.SecantMethod`, `RS.BrentsMethod`). -- `param_set`: Thermodynamics parameter set. -- `T_unsat`: Unsaturated temperature estimate (used for initialization or lower bound) [K]. -- `T_ice`: Temperature with all water as ice (used for upper bound) [K]. -- `T_guess`: Optional user-provided initial guess [K] (or `nothing`). - -# Returns -- Instantiated solver method ready for `RootSolvers.find_zero`. - -# Notes -- For Newton-type methods (`NewtonsMethod`, `NewtonsMethodAD`): Uses `T_guess` if provided, - otherwise `max(T_init_min, T_unsat)`. -- For bracket methods (`SecantMethod`, `BrentsMethod`): Constructs bracket `[T_lo, T_hi]` - where `T_hi` is bounded by `T_ice` and `T_max`. -""" -@inline function _make_sa_solver( - ::Type{RS.NewtonsMethod}, - param_set::APS, - T_unsat, - T_ice, - T_guess, -) - T_init_min = TP.T_init_min(param_set) - T_init = T_guess isa Nothing ? max(T_init_min, T_unsat) : T_guess - return RS.NewtonsMethod(T_init) -end - -@inline function _make_sa_solver( - ::Type{RS.NewtonsMethodAD}, - param_set::APS, - T_unsat, - T_ice, - T_guess, -) - T_init_min = TP.T_init_min(param_set) - T_init = T_guess isa Nothing ? max(T_init_min, T_unsat) : T_guess - return RS.NewtonsMethodAD(T_init) -end - -@inline function _make_sa_solver( - ::Type{RS.SecantMethod}, - param_set::APS, - T_unsat, - T_ice, - T_guess, -) - T_init_min = TP.T_init_min(param_set) - T_lo = T_guess isa Nothing ? max(T_init_min, T_unsat) : max(T_init_min, T_guess) - T_hi = bound_upper_temperature(param_set, T_lo, T_ice) - return RS.SecantMethod(T_lo, T_hi) -end - -@inline function _make_sa_solver( - ::Type{RS.BrentsMethod}, - param_set::APS, - T_unsat, - T_ice, - T_guess, -) - T_init_min = TP.T_init_min(param_set) - # BrentsMethod requires strict bracketing - ignore T_guess - T_lo = max(T_init_min, T_unsat) - T_hi = bound_upper_temperature(param_set, T_lo, T_ice) - return RS.BrentsMethod(T_lo, T_hi) -end - - """ saturation_adjustment( ::Type{M}, # RS.RootSolvingMethod type @@ -253,7 +113,8 @@ end q_tot::Real, maxiter::Int, tol, - [T_guess::Union{Nothing, Real}] + [T_guess::Union{Nothing, Real} = nothing], + [forced_fixed_iters::Bool = false] ) Compute the saturation equilibrium temperature `T` and phase partition `(q_liq, q_ice)` @@ -268,7 +129,10 @@ given pressure `p`, specific internal energy `e_int`, and total specific humidit - `q_tot`: Total specific humidity [kg/kg]. - `maxiter`: Maximum iterations for the solver [dimensionless integer]. - `tol`: Relative tolerance for the temperature solution (or a `RS.RelativeSolutionTolerance`). -- `T_guess`: Optional initial guess for the temperature [K]. +- `T_guess`: Optional initial guess for the temperature [K]. Defaults to `nothing`. +- `forced_fixed_iters`: Optional boolean to force a fixed number of iterations (`maxiter`) + without checking for convergence. Useful for GPU optimization to avoid branch divergence. + When `true`, `T_guess` and `tol` are ignored. Defaults to `false`. # Returns - `NamedTuple` `(; T, q_liq, q_ice, converged)`: @@ -276,6 +140,9 @@ given pressure `p`, specific internal energy `e_int`, and total specific humidit - `q_liq`: Liquid specific humidity [kg/kg] - `q_ice`: Ice specific humidity [kg/kg] - `converged`: Boolean flag indicating if the solver converged + +# Notes +- **GPU broadcasting**: Pass `forced_fixed_iters` as a positional Bool. """ function saturation_adjustment( ::Type{M}, # RS.AbstractMethod type @@ -287,11 +154,11 @@ function saturation_adjustment( maxiter::Int, tol, T_guess = nothing, + forced_fixed_iters::Bool = false, ) where {M} - e_int_sat_given_p = - T -> - internal_energy_sat(param_set, T, air_density(param_set, T, p, q_tot), q_tot) - + if forced_fixed_iters + return saturation_adjustment_fixed_iters(param_set, pe(), p, e_int, q_tot, maxiter) + end q_sat_func = (param_set, T, q_tot) -> q_vap_saturation(param_set, T, air_density(param_set, T, p, q_tot)) @@ -300,11 +167,11 @@ function saturation_adjustment( (param_set, e_int, q_tot) -> air_temperature(param_set, e_int, q_tot, zero(q_tot), q_tot) - roots_func = T -> e_int_sat_given_p(T) - e_int - temp_unsat_func = (param_set, e_int, q_tot) -> air_temperature(param_set, e_int, q_tot) + roots_func = _make_roots_function(M, param_set, pe(), p, e_int, q_tot) + (T, converged) = _saturation_adjustment_generic( M, param_set, @@ -333,7 +200,8 @@ end q_tot::Real, maxiter::Int, tol, - [T_guess::Union{Nothing, Real}] + [T_guess::Union{Nothing, Real} = nothing], + [forced_fixed_iters::Bool = false] ) Compute the saturation equilibrium temperature `T` and phase partition `(q_liq, q_ice)` @@ -350,7 +218,10 @@ Returns a `NamedTuple` `(; T, q_liq, q_ice, converged)`. - `q_tot`: Total specific humidity [kg/kg]. - `maxiter`: Maximum iterations for the solver [dimensionless integer]. - `tol`: Relative tolerance for the temperature solution (or a `RS.RelativeSolutionTolerance`). -- `T_guess`: Optional initial guess for the temperature [K]. +- `T_guess`: Optional initial guess for the temperature [K]. Defaults to `nothing`. +- `forced_fixed_iters`: Optional boolean to force a fixed number of iterations (`maxiter`) + without checking for convergence. Useful for GPU optimization to avoid branch divergence. + When `true`, `T_guess` and `tol` are ignored. Defaults to `false`. # Returns - `NamedTuple` `(; T, q_liq, q_ice, converged)`: @@ -358,6 +229,9 @@ Returns a `NamedTuple` `(; T, q_liq, q_ice, converged)`. - `q_liq`: Liquid specific humidity [kg/kg] - `q_ice`: Ice specific humidity [kg/kg] - `converged`: Boolean flag indicating if the solver converged + +# Notes +- **GPU broadcasting**: Pass `forced_fixed_iters` as a positional Bool. """ function saturation_adjustment( ::Type{M}, # RS.AbstractMethod type @@ -369,11 +243,11 @@ function saturation_adjustment( maxiter::Int, tol, T_guess = nothing, + forced_fixed_iters::Bool = false, ) where {M} - h_sat_given_p = - T -> - enthalpy_sat(param_set, T, air_density(param_set, T, p, q_tot), q_tot) - + if forced_fixed_iters + return saturation_adjustment_fixed_iters(param_set, ph(), p, h, q_tot, maxiter) + end q_sat_func = (param_set, T, q_tot) -> q_vap_saturation(param_set, T, air_density(param_set, T, p, q_tot)) @@ -386,7 +260,7 @@ function saturation_adjustment( (param_set, h, q_tot) -> air_temperature(param_set, ph(), h, q_tot, zero(q_tot), q_tot) - roots_func = T -> h_sat_given_p(T) - h + roots_func = _make_roots_function(M, param_set, ph(), p, h, q_tot) (T, converged) = _saturation_adjustment_generic( M, @@ -416,7 +290,8 @@ end q_tot::Real, maxiter::Int, tol, - [T_guess::Union{Nothing, Real}] + [T_guess::Union{Nothing, Real} = nothing], + [forced_fixed_iters::Bool = false] ) Compute the saturation equilibrium temperature `T` and phase partition `(q_liq, q_ice)` @@ -431,7 +306,10 @@ given pressure `p`, liquid-ice potential temperature `θ_li`, and total specific - `q_tot`: Total specific humidity [kg/kg]. - `maxiter`: Maximum iterations for the solver [dimensionless integer]. - `tol`: Relative tolerance for the temperature solution (or a `RS.RelativeSolutionTolerance`). -- `T_guess`: Optional initial guess for the temperature [K]. +- `T_guess`: Optional initial guess for the temperature [K]. Defaults to `nothing`. +- `forced_fixed_iters`: Optional boolean to force a fixed number of iterations (`maxiter`) + without checking for convergence. Useful for GPU optimization to avoid branch divergence. + When `true`, `T_guess` and `tol` are ignored. Defaults to `false`. # Returns - `NamedTuple` `(; T, q_liq, q_ice, converged)`: @@ -439,6 +317,9 @@ given pressure `p`, liquid-ice potential temperature `θ_li`, and total specific - `q_liq`: Liquid specific humidity [kg/kg] - `q_ice`: Ice specific humidity [kg/kg] - `converged`: Boolean flag indicating if the solver converged + +# Notes +- **GPU broadcasting**: Pass `forced_fixed_iters` as a positional Bool. """ function saturation_adjustment( ::Type{M}, # RS.AbstractMethod type @@ -450,14 +331,18 @@ function saturation_adjustment( maxiter::Int, tol, T_guess = nothing, + forced_fixed_iters::Bool = false, ) where {M} - θ_li_sat_given_p = - T -> begin - _ρ = air_density(param_set, T, p, q_tot) - (_q_liq, _q_ice) = condensate_partition(param_set, T, _ρ, q_tot) - liquid_ice_pottemp_given_pressure(param_set, T, p, q_tot, _q_liq, _q_ice) - end - + if forced_fixed_iters + return saturation_adjustment_fixed_iters( + param_set, + pθ_li(), + p, + θ_li, + q_tot, + maxiter, + ) + end temp_unsat_func = (param_set, θ_li, q_tot) -> air_temperature(param_set, pθ_li(), p, θ_li, q_tot) @@ -478,7 +363,7 @@ function saturation_adjustment( q_tot, ) - roots_func = T -> θ_li_sat_given_p(T) - θ_li + roots_func = _make_roots_function(M, param_set, pθ_li(), p, θ_li, q_tot) (T, converged) = _saturation_adjustment_generic( M, @@ -509,7 +394,8 @@ end q_tot::Real, maxiter::Int, tol, - [T_guess::Union{Nothing, Real}] + [T_guess::Union{Nothing, Real} = nothing], + [forced_fixed_iters::Bool = false] ) Compute the saturation equilibrium temperature `T` and phase partition `(q_liq, q_ice)` @@ -524,7 +410,10 @@ given density `ρ`, liquid-ice potential temperature `θ_li`, and total specific - `q_tot`: Total specific humidity [kg/kg]. - `maxiter`: Maximum iterations for the solver [dimensionless integer]. - `tol`: Relative tolerance for the temperature solution (or a `RS.RelativeSolutionTolerance`). -- `T_guess`: Optional initial guess for the temperature [K]. +- `T_guess`: Optional initial guess for the temperature [K]. Defaults to `nothing`. +- `forced_fixed_iters`: Optional boolean to force a fixed number of iterations (`maxiter`) + without checking for convergence. Useful for GPU optimization to avoid branch divergence. + When `true`, `T_guess` and `tol` are ignored. Defaults to `false`. # Returns - `NamedTuple` `(; T, q_liq, q_ice, converged)`: @@ -532,6 +421,9 @@ given density `ρ`, liquid-ice potential temperature `θ_li`, and total specific - `q_liq`: Liquid specific humidity [kg/kg] - `q_ice`: Ice specific humidity [kg/kg] - `converged`: Boolean flag indicating if the solver converged + +# Notes +- **GPU broadcasting**: Pass `forced_fixed_iters` as a positional Bool. """ function saturation_adjustment( ::Type{M}, # RS.AbstractMethod type @@ -543,13 +435,18 @@ function saturation_adjustment( maxiter::Int, tol, T_guess = nothing, + forced_fixed_iters::Bool = false, ) where {M} - θ_li_sat_given_ρ = - T -> begin - (_q_liq, _q_ice) = condensate_partition(param_set, T, ρ, q_tot) - liquid_ice_pottemp(param_set, T, ρ, q_tot, _q_liq, _q_ice) - end - + if forced_fixed_iters + return saturation_adjustment_fixed_iters( + param_set, + ρθ_li(), + ρ, + θ_li, + q_tot, + maxiter, + ) + end temp_unsat_func = (param_set, θ_li, q_tot) -> air_temperature(param_set, ρθ_li(), ρ, θ_li, q_tot) @@ -569,7 +466,7 @@ function saturation_adjustment( q_tot, ) - roots_func = T -> θ_li_sat_given_ρ(T) - θ_li + roots_func = _make_roots_function(M, param_set, ρθ_li(), ρ, θ_li, q_tot) (T, converged) = _saturation_adjustment_generic( M, @@ -599,7 +496,8 @@ end q_tot::Real, maxiter::Int, tol, - [T_guess::Union{Nothing, Real}] + [T_guess::Union{Nothing, Real} = nothing], + [forced_fixed_iters::Bool = false] ) Compute the saturation equilibrium temperature `T` and phase partition `(q_liq, q_ice)` @@ -614,7 +512,10 @@ given pressure `p`, density `ρ`, and total specific humidity `q_tot`. - `q_tot`: Total specific humidity [kg/kg]. - `maxiter`: Maximum iterations for the solver [dimensionless integer]. - `tol`: Relative tolerance for the temperature solution (or a `RS.RelativeSolutionTolerance`). -- `T_guess`: Optional initial guess for the temperature [K]. +- `T_guess`: Optional initial guess for the temperature [K]. Defaults to `nothing`. +- `forced_fixed_iters`: Optional boolean to force a fixed number of iterations (`maxiter`) + without checking for convergence. Useful for GPU optimization to avoid branch divergence. + When `true`, `T_guess` and `tol` are ignored. Defaults to `false`. # Returns - `NamedTuple` `(; T, q_liq, q_ice, converged)`: @@ -622,6 +523,9 @@ given pressure `p`, density `ρ`, and total specific humidity `q_tot`. - `q_liq`: Liquid specific humidity [kg/kg] - `q_ice`: Ice specific humidity [kg/kg] - `converged`: Boolean flag indicating if the solver converged + +# Notes +- **GPU broadcasting**: Pass `forced_fixed_iters` as a positional Bool. """ function saturation_adjustment( ::Type{M}, # RS.AbstractMethod type @@ -633,13 +537,11 @@ function saturation_adjustment( maxiter::Int, tol, T_guess = nothing, + forced_fixed_iters::Bool = false, ) where {M} - pressure_sat_given_ρ = - T -> begin - (_q_liq, _q_ice) = condensate_partition(param_set, T, ρ, q_tot) - air_pressure(param_set, T, ρ, q_tot, _q_liq, _q_ice) - end - + if forced_fixed_iters + return saturation_adjustment_fixed_iters(param_set, pρ(), p, ρ, q_tot, maxiter) + end temp_unsat_func = (param_set, p, q_tot) -> air_temperature(param_set, pρ(), p, ρ, q_tot) @@ -651,9 +553,8 @@ function saturation_adjustment( (param_set, p, q_tot) -> air_temperature(param_set, pρ(), p, ρ, q_tot, zero(q_tot), q_tot) - roots_func = T -> pressure_sat_given_ρ(T) - p + roots_func = _make_roots_function(M, param_set, pρ(), p, ρ, q_tot) - # Using generic helper with p as target thermo_var (T, converged) = _saturation_adjustment_generic( M, param_set, @@ -693,6 +594,10 @@ branch-free execution on GPUs. For typical atmospheric conditions (T < 320 K), this achieves better than 0.1 K accuracy. For more control over solver parameters, use the full signature with explicit method type. + +# Returns +- `NamedTuple` `(; T, q_liq, q_ice)` — note: `converged` is not included; fixed iterations + always succeed by construction. """ function saturation_adjustment( param_set::APS, @@ -702,13 +607,8 @@ function saturation_adjustment( q_tot; maxiter::Int = 2, ) - return saturation_adjustment_ρe_fixed_iters( - param_set, - ρ, - e_int, - q_tot, - maxiter, - ) + sa_result = saturation_adjustment_fixed_iters(param_set, ρe(), ρ, e_int, q_tot, maxiter) + return (; sa_result.T, sa_result.q_liq, sa_result.q_ice) end """ @@ -718,16 +618,20 @@ end p, e_int, q_tot; - maxiter::Int = 15, - tol = nothing, - T_guess = nothing, + maxiter::Int = 2, ) -Convenience method for `pe` formulation with reasonable defaults. +Convenience method for `pe` formulation with reasonable GPU-optimized defaults. -Uses `RS.SecantMethod` with `maxiter=15` and relative tolerance `1e-4`. +Uses `RS.NewtonsMethod` with `forced_fixed_iters=true` and `maxiter=2` for fast, +branch-free execution on GPUs. For typical atmospheric conditions (T < 320 K), +this achieves better than 0.1 K accuracy. For more control over solver parameters, use the full signature with explicit method type. + +# Returns +- `NamedTuple` `(; T, q_liq, q_ice)` — note: `converged` is not included; fixed iterations + always succeed by construction. """ function saturation_adjustment( param_set::APS, @@ -735,24 +639,10 @@ function saturation_adjustment( p, e_int, q_tot; - maxiter::Int = 15, - tol = nothing, - T_guess = nothing, + maxiter::Int = 2, ) - FT = eltype(param_set) - _tol = tol === nothing ? RS.RelativeSolutionTolerance(FT(1e-4)) : tol - - return saturation_adjustment( - RS.SecantMethod, - param_set, - pe(), - p, - e_int, - q_tot, - maxiter, - _tol, - T_guess, - ) + sa_result = saturation_adjustment_fixed_iters(param_set, pe(), p, e_int, q_tot, maxiter) + return (; sa_result.T, sa_result.q_liq, sa_result.q_ice) end """ @@ -762,16 +652,20 @@ end p, h, q_tot; - maxiter::Int = 15, - tol = nothing, - T_guess = nothing, + maxiter::Int = 2, ) -Convenience method for `ph` formulation with reasonable defaults. +Convenience method for `ph` formulation with reasonable GPU-optimized defaults. -Uses `RS.SecantMethod` with `maxiter=15` and relative tolerance `1e-4`. +Uses `RS.NewtonsMethod` with `forced_fixed_iters=true` and `maxiter=2` for fast, +branch-free execution on GPUs. For typical atmospheric conditions (T < 320 K), +this achieves better than 0.1 K accuracy. For more control over solver parameters, use the full signature with explicit method type. + +# Returns +- `NamedTuple` `(; T, q_liq, q_ice)` — note: `converged` is not included; fixed iterations + always succeed by construction. """ function saturation_adjustment( param_set::APS, @@ -779,24 +673,11 @@ function saturation_adjustment( p, h, q_tot; - maxiter::Int = 15, + maxiter::Int = 2, tol = nothing, - T_guess = nothing, ) - FT = eltype(param_set) - _tol = tol === nothing ? RS.RelativeSolutionTolerance(FT(1e-4)) : tol - - return saturation_adjustment( - RS.SecantMethod, - param_set, - ph(), - p, - h, - q_tot, - maxiter, - _tol, - T_guess, - ) + sa_result = saturation_adjustment_fixed_iters(param_set, ph(), p, h, q_tot, maxiter) + return (; sa_result.T, sa_result.q_liq, sa_result.q_ice) end """ @@ -806,16 +687,20 @@ end p, θ_li, q_tot; - maxiter::Int = 15, - tol = nothing, - T_guess = nothing, + maxiter::Int = 2, ) -Convenience method for `pθ_li` formulation with reasonable defaults. +Convenience method for `pθ_li` formulation with reasonable GPU-optimized defaults. -Uses `RS.SecantMethod` with `maxiter=15` and relative tolerance `1e-4`. +Uses `RS.NewtonsMethod` with `forced_fixed_iters=true` and `maxiter=2` for fast, +branch-free execution on GPUs. For typical atmospheric conditions (T < 320 K), +this achieves better than 0.1 K accuracy. For more control over solver parameters, use the full signature with explicit method type. + +# Returns +- `NamedTuple` `(; T, q_liq, q_ice)` — note: `converged` is not included; fixed iterations + always succeed by construction. """ function saturation_adjustment( param_set::APS, @@ -823,24 +708,12 @@ function saturation_adjustment( p, θ_li, q_tot; - maxiter::Int = 15, + maxiter::Int = 2, tol = nothing, - T_guess = nothing, ) - FT = eltype(param_set) - _tol = tol === nothing ? RS.RelativeSolutionTolerance(FT(1e-4)) : tol - - return saturation_adjustment( - RS.SecantMethod, - param_set, - pθ_li(), - p, - θ_li, - q_tot, - maxiter, - _tol, - T_guess, - ) + sa_result = + saturation_adjustment_fixed_iters(param_set, pθ_li(), p, θ_li, q_tot, maxiter) + return (; sa_result.T, sa_result.q_liq, sa_result.q_ice) end """ @@ -850,16 +723,20 @@ end ρ, θ_li, q_tot; - maxiter::Int = 15, - tol = nothing, - T_guess = nothing, + maxiter::Int = 2, ) -Convenience method for `ρθ_li` formulation with reasonable defaults. +Convenience method for `ρθ_li` formulation with reasonable GPU-optimized defaults. -Uses `RS.SecantMethod` with `maxiter=15` and relative tolerance `1e-4`. +Uses `RS.NewtonsMethod` with `forced_fixed_iters=true` and `maxiter=2` for fast, +branch-free execution on GPUs. For typical atmospheric conditions (T < 320 K), +this achieves better than 0.1 K accuracy. For more control over solver parameters, use the full signature with explicit method type. + +# Returns +- `NamedTuple` `(; T, q_liq, q_ice)` — note: `converged` is not included; fixed iterations + always succeed by construction. """ function saturation_adjustment( param_set::APS, @@ -867,24 +744,11 @@ function saturation_adjustment( ρ, θ_li, q_tot; - maxiter::Int = 15, - tol = nothing, - T_guess = nothing, + maxiter::Int = 2, ) - FT = eltype(param_set) - _tol = tol === nothing ? RS.RelativeSolutionTolerance(FT(1e-4)) : tol - - return saturation_adjustment( - RS.SecantMethod, - param_set, - ρθ_li(), - ρ, - θ_li, - q_tot, - maxiter, - _tol, - T_guess, - ) + sa_result = + saturation_adjustment_fixed_iters(param_set, ρθ_li(), ρ, θ_li, q_tot, maxiter) + return (; sa_result.T, sa_result.q_liq, sa_result.q_ice) end """ @@ -894,16 +758,20 @@ end p, ρ, q_tot; - maxiter::Int = 15, - tol = nothing, - T_guess = nothing, + maxiter::Int = 2, ) -Convenience method for `pρ` formulation with reasonable defaults. +Convenience method for `pρ` formulation with reasonable GPU-optimized defaults. -Uses `RS.SecantMethod` with `maxiter=15` and relative tolerance `1e-4`. +Uses `RS.NewtonsMethod` with `forced_fixed_iters=true` and `maxiter=2` for fast, +branch-free execution on GPUs. For typical atmospheric conditions (T < 320 K), +this achieves better than 0.1 K accuracy. For more control over solver parameters, use the full signature with explicit method type. + +# Returns +- `NamedTuple` `(; T, q_liq, q_ice)` — note: `converged` is not included; fixed iterations + always succeed by construction. """ function saturation_adjustment( param_set::APS, @@ -911,31 +779,277 @@ function saturation_adjustment( p, ρ, q_tot; - maxiter::Int = 15, - tol = nothing, - T_guess = nothing, + maxiter::Int = 2, ) - FT = eltype(param_set) - _tol = tol === nothing ? RS.RelativeSolutionTolerance(FT(1e-4)) : tol + sa_result = saturation_adjustment_fixed_iters(param_set, pρ(), p, ρ, q_tot, maxiter) + return (; sa_result.T, sa_result.q_liq, sa_result.q_ice) +end - return saturation_adjustment( - RS.SecantMethod, - param_set, - pρ(), - p, - ρ, - q_tot, - maxiter, - _tol, - T_guess, - ) +# --------------------------------------------- +# GPU-optimized: fixed iteration methods +# --------------------------------------------- + +""" + saturation_adjustment_fixed_iters(param_set, ::ThermoType, args..., maxiter) + +GPU-optimized saturation adjustment using a fixed number of Newton iterations. + +Bypasses standard solver logic (bracketing, unsaturated checks, convergence testing) +to avoid branch divergence on GPUs. Dispatches on thermodynamic formulation type. + +# Supported formulations +- `ρe`: `saturation_adjustment_fixed_iters(param_set, ρe(), ρ, e_int, q_tot, maxiter)` +- `pe`: `saturation_adjustment_fixed_iters(param_set, pe(), p, e_int, q_tot, maxiter)` +- `ph`: `saturation_adjustment_fixed_iters(param_set, ph(), p, h, q_tot, maxiter)` +- `pθ_li`: `saturation_adjustment_fixed_iters(param_set, pθ_li(), p, θ_li, q_tot, maxiter)` +- `ρθ_li`: `saturation_adjustment_fixed_iters(param_set, ρθ_li(), ρ, θ_li, q_tot, maxiter)` +- `pρ`: `saturation_adjustment_fixed_iters(param_set, pρ(), p, ρ, q_tot, maxiter)` + +# Returns +- `NamedTuple` `(; T, q_liq, q_ice, converged)` + +# Notes +- No convergence check is performed; `converged` is always `true`. +- With `maxiter = 2`, temperature accuracy is better than 0.1 K for typical atmospheric + conditions (T < 320 K). +- This is an internal helper function. For the public API, use [`saturation_adjustment`](@ref) + with `forced_fixed_iters=true` as a positional argument. +""" +function saturation_adjustment_fixed_iters end + +@inline function saturation_adjustment_fixed_iters( + param_set::APS, + ::ρe, + ρ, + e_int, + q_tot, + maxiter, +) + T_unsat = air_temperature(param_set, e_int, q_tot) + T_init_min = TP.T_init_min(param_set) + T = max(T_init_min, T_unsat) + + @fastmath for _ in 1:maxiter + e_val = internal_energy_sat(param_set, T, ρ, q_tot) + de_int_dT = ∂e_int_∂T_sat_ρ(param_set, T, ρ, q_tot) + T += (e_int - e_val) / de_int_dT + end + + (q_liq, q_ice) = condensate_partition(param_set, T, ρ, q_tot) + return (; T, q_liq, q_ice, converged = true) end +@inline function saturation_adjustment_fixed_iters( + param_set::APS, + ::pe, + p, + e_int, + q_tot, + maxiter, +) + T_unsat = air_temperature(param_set, e_int, q_tot) + T_init_min = TP.T_init_min(param_set) + T = max(T_init_min, T_unsat) + + @fastmath for _ in 1:maxiter + ρ = air_density(param_set, T, p, q_tot) + e_val = internal_energy_sat(param_set, T, ρ, q_tot) + de_int_dT = ∂e_int_∂T_sat_p(param_set, T, p, q_tot) + T -= (e_val - e_int) / de_int_dT + end + + (ρ, q_liq, q_ice) = _phase_partition_from_T_p(param_set, T, p, q_tot) + return (; T, q_liq, q_ice, converged = true) +end + +@inline function saturation_adjustment_fixed_iters( + param_set::APS, + ::ph, + p, + h, + q_tot, + maxiter, +) + T_unsat = air_temperature(param_set, ph(), h, q_tot, 0, 0) + T_init_min = TP.T_init_min(param_set) + T = max(T_init_min, T_unsat) + + @fastmath for _ in 1:maxiter + ρ = air_density(param_set, T, p, q_tot) + h_val = enthalpy_sat(param_set, T, ρ, q_tot) + dh_dT = ∂h_∂T_sat_p(param_set, T, p, q_tot) + T -= (h_val - h) / dh_dT + end + + (ρ, q_liq, q_ice) = _phase_partition_from_T_p(param_set, T, p, q_tot) + return (; T, q_liq, q_ice, converged = true) +end + +@inline function saturation_adjustment_fixed_iters( + param_set::APS, + ::pθ_li, + p, + θ_li, + q_tot, + maxiter, +) + T_unsat = air_temperature(param_set, pθ_li(), p, θ_li, q_tot) + T_init_min = TP.T_init_min(param_set) + T = max(T_init_min, T_unsat) + + @fastmath for _ in 1:maxiter + ρ = air_density(param_set, T, p, q_tot) + (q_liq, q_ice) = condensate_partition(param_set, T, ρ, q_tot) + θ_li_val = liquid_ice_pottemp_given_pressure(param_set, T, p, q_tot, q_liq, q_ice) + dθ_li_dT = ∂θ_li_∂T_sat_p(param_set, T, p, q_tot) + T -= (θ_li_val - θ_li) / dθ_li_dT + end + + (ρ, q_liq, q_ice) = _phase_partition_from_T_p(param_set, T, p, q_tot) + return (; T, q_liq, q_ice, converged = true) +end + +@inline function saturation_adjustment_fixed_iters( + param_set::APS, + ::ρθ_li, + ρ, + θ_li, + q_tot, + maxiter, +) + T_unsat = air_temperature(param_set, ρθ_li(), ρ, θ_li, q_tot) + T_init_min = TP.T_init_min(param_set) + T = max(T_init_min, T_unsat) + + @fastmath for _ in 1:maxiter + (q_liq, q_ice) = condensate_partition(param_set, T, ρ, q_tot) + θ_li_val = liquid_ice_pottemp(param_set, T, ρ, q_tot, q_liq, q_ice) + dθ_li_dT = ∂θ_li_∂T_sat_ρ(param_set, T, ρ, q_tot) + T -= (θ_li_val - θ_li) / dθ_li_dT + end + + (q_liq, q_ice) = condensate_partition(param_set, T, ρ, q_tot) + return (; T, q_liq, q_ice, converged = true) +end + +@inline function saturation_adjustment_fixed_iters( + param_set::APS, + ::pρ, + p, + ρ, + q_tot, + maxiter, +) + T_unsat = air_temperature(param_set, pρ(), p, ρ, q_tot) + T_init_min = TP.T_init_min(param_set) + T = max(T_init_min, T_unsat) + + @fastmath for _ in 1:maxiter + (q_liq, q_ice) = condensate_partition(param_set, T, ρ, q_tot) + p_val = air_pressure(param_set, T, ρ, q_tot, q_liq, q_ice) + dp_dT = ∂p_∂T_sat_ρ(param_set, T, ρ, q_tot) + T -= (p_val - p) / dp_dT + end + + (q_liq, q_ice) = condensate_partition(param_set, T, ρ, q_tot) + return (; T, q_liq, q_ice, converged = true) +end # --------------------------------------------- -# Helper functions +# Internal helpers # --------------------------------------------- +""" + bound_upper_temperature(param_set, T_lo, T_hi) + +Internal function. Bounds the upper temperature guess `T_hi` for bracket methods. + +Returns `T_hi` bounded by `T_max`, ensuring it is at least `T_lo + 0.1` for +valid numerical initialization. +""" +@inline function bound_upper_temperature(param_set, T_lo, T_hi) + FT = eltype(param_set) + T_max = TP.T_max(param_set) + # Ensure T_hi is physically valid (<= T_max) + T_hi_phys = min(T_max, T_hi) + # Ensure T_hi > T_lo for numerical initialization (use relative tolerance) + return max(T_lo * (1 + FT(1e-3)), T_hi_phys) +end + + +""" + _make_sa_solver(::Type{M}, param_set, T_unsat, T_ice, T_guess) + +Internal helper to construct a root-solving method instance for saturation adjustment. + +# Arguments +- `M`: Root-solving method type (e.g., `RS.NewtonsMethod`, `RS.SecantMethod`, `RS.BrentsMethod`). +- `param_set`: Thermodynamics parameter set. +- `T_unsat`: Unsaturated temperature estimate (used for initialization or lower bound) [K]. +- `T_ice`: Temperature with all water as ice (used for upper bound) [K]. +- `T_guess`: Optional user-provided initial guess [K] (or `nothing`). + +# Returns +- Instantiated solver method ready for `RootSolvers.find_zero`. + +# Notes +- For Newton-type methods (`NewtonsMethod`, `NewtonsMethodAD`): Uses `T_guess` if provided, + otherwise `max(T_init_min, T_unsat)`. +- For bracket methods (`SecantMethod`, `BrentsMethod`): Constructs bracket `[T_lo, T_hi]` + where `T_hi` is bounded by `T_ice` and `T_max`. +""" +@inline function _make_sa_solver( + ::Type{RS.NewtonsMethod}, + param_set::APS, + T_unsat, + T_ice, + T_guess, +) + T_init_min = TP.T_init_min(param_set) + T_init = T_guess isa Nothing ? max(T_init_min, T_unsat) : T_guess + return RS.NewtonsMethod(T_init) +end + +@inline function _make_sa_solver( + ::Type{RS.NewtonsMethodAD}, + param_set::APS, + T_unsat, + T_ice, + T_guess, +) + T_init_min = TP.T_init_min(param_set) + T_init = T_guess isa Nothing ? max(T_init_min, T_unsat) : T_guess + return RS.NewtonsMethodAD(T_init) +end + +@inline function _make_sa_solver( + ::Type{RS.SecantMethod}, + param_set::APS, + T_unsat, + T_ice, + T_guess, +) + T_init_min = TP.T_init_min(param_set) + T_lo = T_guess isa Nothing ? max(T_init_min, T_unsat) : max(T_init_min, T_guess) + T_hi = bound_upper_temperature(param_set, T_lo, T_ice) + return RS.SecantMethod(T_lo, T_hi) +end + +@inline function _make_sa_solver( + ::Type{RS.BrentsMethod}, + param_set::APS, + T_unsat, + T_ice, + T_guess, +) + T_init_min = TP.T_init_min(param_set) + # BrentsMethod requires strict bracketing - ignore T_guess + T_lo = max(T_init_min, T_unsat) + T_hi = bound_upper_temperature(param_set, T_lo, T_ice) + return RS.BrentsMethod(T_lo, T_hi) +end + + """ internal_energy_sat(param_set, T, ρ, q_tot) @@ -983,14 +1097,20 @@ excess (see [`saturation_excess`](@ref)). end """ - _make_roots_function(::Type{M}, param_set, ρ, e_int, q_tot) + _make_roots_function(::Type{M}, param_set, ::ThermoType, args..., q_tot) Helper function to create the root function for Newton's method (with derivative) or -other methods (without derivative), dispatching on the method type. +other methods (without derivative), dispatching on the method type and thermo type. + +Returns `(f, f')` for `NewtonsMethod`, or just `f` for other methods. """ +function _make_roots_function end + +# ρe formulation @inline function _make_roots_function( ::Type{RS.NewtonsMethod}, param_set::APS, + ::ρe, ρ, e_int, q_tot, @@ -998,13 +1118,14 @@ other methods (without derivative), dispatching on the method type. return _T -> begin T_val = ReLU(_T) f = e_int - internal_energy_sat(param_set, T_val, ρ, q_tot) - (f, -∂e_int_∂T_sat(param_set, T_val, ρ, q_tot)) + (f, -∂e_int_∂T_sat_ρ(param_set, T_val, ρ, q_tot)) end end @inline function _make_roots_function( ::Type{M}, param_set::APS, + ::ρe, ρ, e_int, q_tot, @@ -1015,6 +1136,164 @@ end end end +# pe formulation +@inline function _make_roots_function( + ::Type{RS.NewtonsMethod}, + param_set::APS, + ::pe, + p, + e_int, + q_tot, +) + return _T -> begin + T_val = ReLU(_T) + _ρ = air_density(param_set, T_val, p, q_tot) + f = internal_energy_sat(param_set, T_val, _ρ, q_tot) - e_int + (f, ∂e_int_∂T_sat_p(param_set, T_val, p, q_tot)) + end +end + +@inline function _make_roots_function( + ::Type{M}, + param_set::APS, + ::pe, + p, + e_int, + q_tot, +) where {M} + return T -> begin + _ρ = air_density(param_set, T, p, q_tot) + internal_energy_sat(param_set, T, _ρ, q_tot) - e_int + end +end + +# ph formulation +@inline function _make_roots_function( + ::Type{RS.NewtonsMethod}, + param_set::APS, + ::ph, + p, + h, + q_tot, +) + return _T -> begin + T_val = ReLU(_T) + _ρ = air_density(param_set, T_val, p, q_tot) + f = enthalpy_sat(param_set, T_val, _ρ, q_tot) - h + (f, ∂h_∂T_sat_p(param_set, T_val, p, q_tot)) + end +end + +@inline function _make_roots_function( + ::Type{M}, + param_set::APS, + ::ph, + p, + h, + q_tot, +) where {M} + return T -> begin + _ρ = air_density(param_set, T, p, q_tot) + enthalpy_sat(param_set, T, _ρ, q_tot) - h + end +end + +# pθ_li formulation +@inline function _make_roots_function( + ::Type{RS.NewtonsMethod}, + param_set::APS, + ::pθ_li, + p, + θ_li, + q_tot, +) + return _T -> begin + T_val = ReLU(_T) + _ρ = air_density(param_set, T_val, p, q_tot) + (_q_liq, _q_ice) = condensate_partition(param_set, T_val, _ρ, q_tot) + f = + liquid_ice_pottemp_given_pressure(param_set, T_val, p, q_tot, _q_liq, _q_ice) - θ_li + (f, ∂θ_li_∂T_sat_p(param_set, T_val, p, q_tot)) + end +end + +@inline function _make_roots_function( + ::Type{M}, + param_set::APS, + ::pθ_li, + p, + θ_li, + q_tot, +) where {M} + return T -> begin + _ρ = air_density(param_set, T, p, q_tot) + (_q_liq, _q_ice) = condensate_partition(param_set, T, _ρ, q_tot) + liquid_ice_pottemp_given_pressure(param_set, T, p, q_tot, _q_liq, _q_ice) - θ_li + end +end + +# ρθ_li formulation +@inline function _make_roots_function( + ::Type{RS.NewtonsMethod}, + param_set::APS, + ::ρθ_li, + ρ, + θ_li, + q_tot, +) + return _T -> begin + T_val = ReLU(_T) + (_q_liq, _q_ice) = condensate_partition(param_set, T_val, ρ, q_tot) + f = liquid_ice_pottemp(param_set, T_val, ρ, q_tot, _q_liq, _q_ice) - θ_li + (f, ∂θ_li_∂T_sat_ρ(param_set, T_val, ρ, q_tot)) + end +end + +@inline function _make_roots_function( + ::Type{M}, + param_set::APS, + ::ρθ_li, + ρ, + θ_li, + q_tot, +) where {M} + return T -> begin + (_q_liq, _q_ice) = condensate_partition(param_set, T, ρ, q_tot) + liquid_ice_pottemp(param_set, T, ρ, q_tot, _q_liq, _q_ice) - θ_li + end +end + +# pρ formulation +@inline function _make_roots_function( + ::Type{RS.NewtonsMethod}, + param_set::APS, + ::pρ, + p, + ρ, + q_tot, +) + return _T -> begin + T_val = ReLU(_T) + (_q_liq, _q_ice) = condensate_partition(param_set, T_val, ρ, q_tot) + f = air_pressure(param_set, T_val, ρ, q_tot, _q_liq, _q_ice) - p + (f, ∂p_∂T_sat_ρ(param_set, T_val, ρ, q_tot)) + end +end + +@inline function _make_roots_function( + ::Type{M}, + param_set::APS, + ::pρ, + p, + ρ, + q_tot, +) where {M} + return T -> begin + (_q_liq, _q_ice) = condensate_partition(param_set, T, ρ, q_tot) + air_pressure(param_set, T, ρ, q_tot, _q_liq, _q_ice) - p + end +end + """ _phase_partition_from_T_p(param_set, T, p, q_tot) @@ -1138,80 +1417,313 @@ end # ------------------------------- """ - ∂e_int_∂T(param_set, T, e_int, ρ, q_tot, q_vap_sat) + _saturation_derivative_vars(param_set, T, ρ, q_tot, q_vap_sat, ::Val{:ρ}) + _saturation_derivative_vars(param_set, T, ρ, q_tot, q_vap_sat, ::Val{:p}) -Derivative of internal energy with respect to temperature at saturation. +Helper to compute common intermediate variables for saturation derivatives. -# Arguments -- `param_set`: Parameter set. -- `T`: Temperature [K]. -- `e_int`: Internal energy [J/kg] (unused, for interface consistency). -- `ρ`: Density [kg/m³]. -- `q_tot`: Total specific humidity [kg/kg]. -- `q_vap_sat`: (Optional) Saturation vapor specific humidity [kg/kg]. +Returns a named tuple with phase partition and derivative information: +- `λ`, `q_c`, `q_liq`, `q_ice`: Phase partition variables +- `∂λ_∂T`, `∂qvs_∂T`: Temperature derivatives of liquid fraction and saturation humidity +- `∂q_liq_∂T`, `∂q_ice_∂T`: Phase partition derivatives + +The `Val{:ρ}` version computes `∂qvs_∂T` at fixed density (includes `-1/T` term). +The `Val{:p}` version computes `∂qvs_∂T` at fixed pressure (no `-1/T` term). """ -@inline function ∂e_int_∂T( +@inline function _saturation_derivative_vars( param_set::APS, T, - e_int, ρ, q_tot, - q_vap_sat = q_vap_saturation(param_set, T, ρ), + q_vap_sat, + ::Val{:ρ}, ) - FT = eltype(param_set) - (q_liq, q_ice) = condensate_partition(param_set, T, ρ, q_tot) - q_cond = q_liq + q_ice + λ = liquid_fraction_ramp(param_set, T) + q_c = saturation_excess(param_set, T, ρ, q_tot) + q_liq = λ * q_c + q_ice = (1 - λ) * q_c # ∂λ/∂T for λ = ((T - Tⁱ) / (Tᶠ - Tⁱ))^n Tᶠ = TP.T_freeze(param_set) Tⁱ = TP.T_icenuc(param_set) n = TP.pow_icenuc(param_set) - ∂λ_∂T = ifelse( - Tⁱ < T < Tᶠ, - n / (Tᶠ - Tⁱ) * ((T - Tⁱ) / (Tᶠ - Tⁱ))^(n - 1), - zero(T), - ) + ∂λ_∂T = ifelse(Tⁱ < T < Tᶠ, n / (Tᶠ - Tⁱ) * ((T - Tⁱ) / (Tᶠ - Tⁱ))^(n - 1), zero(T)) - # ∂q_vap_sat/∂T from Clausius-Clapeyron - _∂q_vap_sat_∂T = ∂q_vap_sat_∂T(param_set, T, ρ) + # ∂q_vap_sat/∂T at fixed ρ (includes -1/T term) + ∂qvs_∂T = ∂q_vap_sat_∂T(param_set, T, ρ) - # Derivative of internal energy with respect to temperature - # Note: we need to handle the unsaturated case explicitly because - # `internal_energy_sat` will revert to `cvm * T + (terms independent of T)` - # when q_tot <= q_vap_sat, so that q_vap = q_tot. + # Phase partition derivatives + ∂q_liq_∂T = ∂λ_∂T * q_c + λ * (-∂qvs_∂T) + ∂q_ice_∂T = -∂λ_∂T * q_c + (1 - λ) * (-∂qvs_∂T) - # Unsaturated case: ∂e_int/∂T = cv_d*(1-q_tot) + cv_v*q_tot - cvm_unsat = cv_m(param_set, q_tot, zero(q_tot), zero(q_tot)) + return (; λ, q_c, q_liq, q_ice, ∂λ_∂T, ∂qvs_∂T, ∂q_liq_∂T, ∂q_ice_∂T) +end - # Saturated case - # Derivatives of phase fractions (when saturated, ∂q_c/∂T = -∂q_vap_sat/∂T) - # q_c = q_tot - q_vap_sat +@inline function _saturation_derivative_vars( + param_set::APS, + T, + ρ, + q_tot, + q_vap_sat, + ::Val{:p}, +) λ = liquid_fraction_ramp(param_set, T) - ∂q_c_∂T = -_∂q_vap_sat_∂T - ∂q_liq_∂T = ∂λ_∂T * q_cond + λ * ∂q_c_∂T - ∂q_ice_∂T = -∂λ_∂T * q_cond + (1 - λ) * ∂q_c_∂T - ∂q_vap_∂T = _∂q_vap_sat_∂T + q_c = saturation_excess(param_set, T, ρ, q_tot) + q_liq = λ * q_c + q_ice = (1 - λ) * q_c + + # ∂λ/∂T for λ = ((T - Tⁱ) / (Tᶠ - Tⁱ))^n + Tᶠ = TP.T_freeze(param_set) + Tⁱ = TP.T_icenuc(param_set) + n = TP.pow_icenuc(param_set) + ∂λ_∂T = ifelse(Tⁱ < T < Tᶠ, n / (Tᶠ - Tⁱ) * ((T - Tⁱ) / (Tᶠ - Tⁱ))^(n - 1), zero(T)) + + # ∂q_vap_sat/∂T at fixed p (no -1/T term; cancels with ∂ρ/∂T|_p = -ρ/T) + R_v = TP.R_v(param_set) + L = latent_heat_mixed(param_set, T, λ) + ∂qvs_∂T = q_vap_sat * L / (R_v * T^2) + + # Phase partition derivatives + ∂q_liq_∂T = ∂λ_∂T * q_c + λ * (-∂qvs_∂T) + ∂q_ice_∂T = -∂λ_∂T * q_c + (1 - λ) * (-∂qvs_∂T) + + return (; λ, q_c, q_liq, q_ice, ∂λ_∂T, ∂qvs_∂T, ∂q_liq_∂T, ∂q_ice_∂T) +end + +""" + ∂e_int_∂T_sat_ρ(param_set, T, ρ, q_tot) + +Derivative of `internal_energy_sat` with respect to temperature at fixed density. + +Uses `∂q_vap_sat/∂T|_ρ` from Clausius-Clapeyron, which includes the `-1/T` term +from the density dependence of saturation vapor pressure. +""" +@inline function ∂e_int_∂T_sat_ρ(param_set::APS, T, ρ, q_tot) + q_vap_sat = q_vap_saturation(param_set, T, ρ) + cvm_unsat = cv_m(param_set, q_tot, zero(q_tot), zero(q_tot)) + + vars = _saturation_derivative_vars(param_set, T, ρ, q_tot, q_vap_sat, Val(:ρ)) + + # Component internal energies + e_vap = internal_energy_vapor(param_set, T) + e_liq = internal_energy_liquid(param_set, T) + e_ice = internal_energy_ice(param_set, T) + + # Full derivative: cv_m + Σ(e_i * ∂q_i/∂T) + cvm_sat = cv_m(param_set, q_tot, vars.q_liq, vars.q_ice) + de_dT_sat = + cvm_sat + e_vap * vars.∂qvs_∂T + e_liq * vars.∂q_liq_∂T + e_ice * vars.∂q_ice_∂T + + return ifelse(q_tot <= q_vap_sat, cvm_unsat, de_dT_sat) +end + +""" + ∂e_int_∂T_sat_p(param_set, T, p, q_tot) + +Derivative of `internal_energy_sat` with respect to temperature at fixed pressure. + +Uses `∂q_vap_sat/∂T|_p = q_vap_sat * L / (R_v T²)` (Clausius-Clapeyron at constant +pressure), which differs from the constant-density form used in [`∂e_int_∂T_sat_ρ`](@ref) +by an additional `q_vap_sat / T` term. +""" +@inline function ∂e_int_∂T_sat_p(param_set::APS, T, p, q_tot) + ρ = air_density(param_set, T, p, q_tot) + q_vap_sat = q_vap_saturation(param_set, T, ρ) + cvm_unsat = cv_m(param_set, q_tot, zero(q_tot), zero(q_tot)) + + vars = _saturation_derivative_vars(param_set, T, ρ, q_tot, q_vap_sat, Val(:p)) # Component internal energies e_vap = internal_energy_vapor(param_set, T) e_liq = internal_energy_liquid(param_set, T) e_ice = internal_energy_ice(param_set, T) - # Full derivative: ∂e_int/∂T = cv_m + Σ(e_i * ∂q_i/∂T) - cvm_sat = cv_m(param_set, q_tot, q_liq, q_ice) - de_int_dT_sat = cvm_sat + e_vap * ∂q_vap_∂T + e_liq * ∂q_liq_∂T + e_ice * ∂q_ice_∂T + # Full derivative: cv_m + Σ(e_i * ∂q_i/∂T) + cvm_sat = cv_m(param_set, q_tot, vars.q_liq, vars.q_ice) + de_dT_sat = + cvm_sat + e_vap * vars.∂qvs_∂T + e_liq * vars.∂q_liq_∂T + e_ice * vars.∂q_ice_∂T - return ifelse(q_tot <= q_vap_sat, cvm_unsat, de_int_dT_sat) + return ifelse(q_tot <= q_vap_sat, cvm_unsat, de_dT_sat) end """ - ∂e_int_∂T_sat(param_set, T, ρ, q_tot) + ∂h_∂T_sat_p(param_set, T, p, q_tot) + +Derivative of `enthalpy_sat` with respect to temperature at fixed pressure. -Helper to compute `∂e_int/∂T` at saturation for Newton's method, -calculating necessary intermediate variables. +Structured identically to [`∂e_int_∂T_sat_p`](@ref) but with component enthalpies +(`cp_m` instead of `cv_m`, `enthalpy_vapor` instead of `internal_energy_vapor`, etc.). """ -@inline function ∂e_int_∂T_sat(param_set::APS, T, ρ, q_tot) +@inline function ∂h_∂T_sat_p(param_set::APS, T, p, q_tot) + ρ = air_density(param_set, T, p, q_tot) q_vap_sat = q_vap_saturation(param_set, T, ρ) - e_int_sat = internal_energy_sat(param_set, T, ρ, q_tot) - return ∂e_int_∂T(param_set, T, e_int_sat, ρ, q_tot, q_vap_sat) + cpm_unsat = cp_m(param_set, q_tot, zero(q_tot), zero(q_tot)) + + vars = _saturation_derivative_vars(param_set, T, ρ, q_tot, q_vap_sat, Val(:p)) + + # Component enthalpies + h_vap = enthalpy_vapor(param_set, T) + h_liq = enthalpy_liquid(param_set, T) + h_ice = enthalpy_ice(param_set, T) + + # Full derivative: cp_m + Σ(h_i * ∂q_i/∂T) + cpm_sat = cp_m(param_set, q_tot, vars.q_liq, vars.q_ice) + dh_dT_sat = + cpm_sat + h_vap * vars.∂qvs_∂T + h_liq * vars.∂q_liq_∂T + h_ice * vars.∂q_ice_∂T + + return ifelse(q_tot <= q_vap_sat, cpm_unsat, dh_dT_sat) +end + +""" + ∂θ_li_∂T_sat_p(param_set, T, p, q_tot) + +Derivative of `liquid_ice_pottemp_given_pressure` (at saturation equilibrium) +with respect to temperature at fixed pressure. + +Uses the product rule on `θ_li = θ * (1 - L_c / (cp_m T))`, differentiating +each factor through the T-dependent phase partition. +""" +@inline function ∂θ_li_∂T_sat_p(param_set::APS, T, p, q_tot) + ρ = air_density(param_set, T, p, q_tot) + q_vap_sat = q_vap_saturation(param_set, T, ρ) + + vars = _saturation_derivative_vars(param_set, T, ρ, q_tot, q_vap_sat, Val(:p)) + + # Current thermodynamic state + R_m = gas_constant_air(param_set, q_tot, vars.q_liq, vars.q_ice) + _cp_m = cp_m(param_set, q_tot, vars.q_liq, vars.q_ice) + p0 = TP.p_ref_theta(param_set) + α = R_m / _cp_m + Π = fast_power(p / p0, α) + θ = T / Π + L_c = humidity_weighted_latent_heat(param_set, vars.q_liq, vars.q_ice) + F = 1 - L_c / (_cp_m * T) + + # Unsaturated: α is constant, L_c = 0, F = 1, so ∂θ_li/∂T = θ/T + dθ_li_dT_unsat = θ / T + + # Parameters + R_v = TP.R_v(param_set) + cp_v = TP.cp_v(param_set) + cp_l = TP.cp_l(param_set) + cp_i = TP.cp_i(param_set) + LH_v0 = TP.LH_v0(param_set) + LH_s0 = TP.LH_s0(param_set) + + # ∂R_m/∂T and ∂cp_m/∂T + ∂R_m_∂T = R_v * vars.∂qvs_∂T + ∂cp_m_∂T = (cp_l - cp_v) * vars.∂q_liq_∂T + (cp_i - cp_v) * vars.∂q_ice_∂T + + # ∂α/∂T = (∂R_m * cp_m - R_m * ∂cp_m) / cp_m² + ∂α_∂T = (∂R_m_∂T * _cp_m - R_m * ∂cp_m_∂T) / _cp_m^2 + + # ∂θ/∂T = θ * (1/T + ln(p₀/p) * ∂α/∂T) + ∂θ_∂T = θ * (1 / T + log(p0 / p) * ∂α_∂T) + + # ∂L_c/∂T + ∂L_c_∂T = LH_v0 * vars.∂q_liq_∂T + LH_s0 * vars.∂q_ice_∂T + + # ∂F/∂T = -1/(cp_m*T) * (∂L_c/∂T - L_c * (1/T + ∂cp_m/∂T / cp_m)) + ∂F_∂T = -1 / (_cp_m * T) * (∂L_c_∂T - L_c * (1 / T + ∂cp_m_∂T / _cp_m)) + + # Product rule: ∂(θ * F)/∂T + dθ_li_dT_sat = ∂θ_∂T * F + θ * ∂F_∂T + + return ifelse(q_tot <= q_vap_sat, dθ_li_dT_unsat, dθ_li_dT_sat) +end + +""" + ∂θ_li_∂T_sat_ρ(param_set, T, ρ, q_tot) + +Derivative of `liquid_ice_pottemp` (at saturation equilibrium) +with respect to temperature at fixed density. + +Structured like [`∂θ_li_∂T_sat_p`](@ref), but accounts for `p = ρ R_m T` varying +with `T`. This introduces two corrections to `∂θ/∂T`: + - an extra `-α/T` from `∂ln(p)/∂T|_ρ = 1/T + ∂R_m/∂T / R_m`, and + - an extra `-∂R_m/∂T / cp_m` from the same. + +The `∂q_vap_sat/∂T` also differs: at fixed ρ it carries an extra `-q_vap_sat / T` +relative to the fixed-p Clausius–Clapeyron form. +""" +@inline function ∂θ_li_∂T_sat_ρ(param_set::APS, T, ρ, q_tot) + q_vap_sat = q_vap_saturation(param_set, T, ρ) + + vars = _saturation_derivative_vars(param_set, T, ρ, q_tot, q_vap_sat, Val(:ρ)) + + # Current thermodynamic state + R_m = gas_constant_air(param_set, q_tot, vars.q_liq, vars.q_ice) + _cp_m = cp_m(param_set, q_tot, vars.q_liq, vars.q_ice) + p0 = TP.p_ref_theta(param_set) + p = ρ * R_m * T # ideal gas at fixed ρ + α = R_m / _cp_m + Π = fast_power(p / p0, α) + θ = T / Π + L_c = humidity_weighted_latent_heat(param_set, vars.q_liq, vars.q_ice) + F = 1 - L_c / (_cp_m * T) + + # Unsaturated: q_liq = q_ice = 0, L_c = 0, F = 1, R_m constant + # θ = (p0/(ρ R_m))^α · T^(1-α) => dθ/dT = (1-α)·θ/T + dθ_li_dT_unsat = θ * (1 - α) / T + + # Parameters + R_v = TP.R_v(param_set) + cp_v = TP.cp_v(param_set) + cp_l = TP.cp_l(param_set) + cp_i = TP.cp_i(param_set) + LH_v0 = TP.LH_v0(param_set) + LH_s0 = TP.LH_s0(param_set) + + # ∂R_m/∂T and ∂cp_m/∂T + ∂R_m_∂T = R_v * vars.∂qvs_∂T + ∂cp_m_∂T = (cp_l - cp_v) * vars.∂q_liq_∂T + (cp_i - cp_v) * vars.∂q_ice_∂T + + # ∂α/∂T = (∂R_m · cp_m - R_m · ∂cp_m) / cp_m² + ∂α_∂T = (∂R_m_∂T * _cp_m - R_m * ∂cp_m_∂T) / _cp_m^2 + + # ∂θ/∂T at fixed ρ: extra -α/T and -∂R_m/∂T/cp_m vs. fixed p + ∂θ_∂T = θ * ((1 - α) / T + log(p0 / p) * ∂α_∂T - ∂R_m_∂T / _cp_m) + + # ∂L_c/∂T + ∂L_c_∂T = LH_v0 * vars.∂q_liq_∂T + LH_s0 * vars.∂q_ice_∂T + + # ∂F/∂T = -1/(cp_m·T) · (∂L_c/∂T - L_c · (1/T + ∂cp_m/∂T / cp_m)) + ∂F_∂T = -1 / (_cp_m * T) * (∂L_c_∂T - L_c * (1 / T + ∂cp_m_∂T / _cp_m)) + + # Product rule: ∂(θ · F)/∂T + dθ_li_dT_sat = ∂θ_∂T * F + θ * ∂F_∂T + + return ifelse(q_tot <= q_vap_sat, dθ_li_dT_unsat, dθ_li_dT_sat) +end + +""" + ∂p_∂T_sat_ρ(param_set, T, ρ, q_tot) + +Derivative of `air_pressure` (at saturation equilibrium) with respect to +temperature at fixed density. + +From `p = ρ R_m T` and `∂R_m/∂T = R_v · ∂q_vap_sat/∂T`: + + ∂p/∂T|_ρ = ρ · (R_m + T · ∂R_m/∂T) + +Unsaturated: `R_m` is constant, so `∂p/∂T = ρ R_m`. +""" +@inline function ∂p_∂T_sat_ρ(param_set::APS, T, ρ, q_tot) + R_v = TP.R_v(param_set) + q_vap_sat = q_vap_saturation(param_set, T, ρ) + + # Unsaturated: R_m constant + R_m_unsat = gas_constant_air(param_set, q_tot, zero(q_tot), zero(q_tot)) + dp_dT_unsat = ρ * R_m_unsat + + vars = _saturation_derivative_vars(param_set, T, ρ, q_tot, q_vap_sat, Val(:ρ)) + R_m = gas_constant_air(param_set, q_tot, vars.q_liq, vars.q_ice) + + # ∂R_m/∂T + ∂R_m_∂T = R_v * vars.∂qvs_∂T + + # p = ρ R_m T => ∂p/∂T = ρ(R_m + T · ∂R_m/∂T) + dp_dT_sat = ρ * (R_m + T * ∂R_m_∂T) + + return ifelse(q_tot <= q_vap_sat, dp_dT_unsat, dp_dT_sat) end diff --git a/test/ad_tests.jl b/test/ad_tests.jl index 6a6b7c36..c1bbd3b2 100644 --- a/test/ad_tests.jl +++ b/test/ad_tests.jl @@ -278,9 +278,9 @@ using Thermodynamics: ρe, pe, ph atol = FT(1e-6), ) - # Test ∂e_int_∂T_sat + # Test ∂e_int_∂T_sat_ρ # Analytical derivative (approximate) - de_int_dT_analytical = TD.∂e_int_∂T_sat(param_set, T, ρ, q_tot) + de_int_dT_analytical = TD.∂e_int_∂T_sat_ρ(param_set, T, ρ, q_tot) # ForwardDiff derivative f_e_int_sat(T_) = TD.internal_energy_sat(param_set, T_, ρ, q_tot) @@ -294,6 +294,105 @@ using Thermodynamics: ρe, pe, ph ) end end + + # Test derivatives at fixed pressure + p_range = FT.([80000, 100000, 101325]) + for T in T_range, p in p_range, q_tot in q_tot_vals + @testset "T=$T, p=$p, q_tot=$q_tot" begin + # Test ∂e_int_∂T_sat_p + de_int_dT_p_analytical = TD.∂e_int_∂T_sat_p(param_set, T, p, q_tot) + + f_e_int_sat_p(T_) = begin + ρ_ = TD.air_density(param_set, T_, p, q_tot) + TD.internal_energy_sat(param_set, T_, ρ_, q_tot) + end + de_int_dT_p_AD = ForwardDiff.derivative(f_e_int_sat_p, T) + + @test isapprox( + de_int_dT_p_analytical, + de_int_dT_p_AD; + rtol = FT(5e-2), + atol = FT(1e-6), + ) + + # Test ∂h_∂T_sat_p + dh_dT_p_analytical = TD.∂h_∂T_sat_p(param_set, T, p, q_tot) + + f_h_sat_p(T_) = begin + ρ_ = TD.air_density(param_set, T_, p, q_tot) + TD.enthalpy_sat(param_set, T_, ρ_, q_tot) + end + dh_dT_p_AD = ForwardDiff.derivative(f_h_sat_p, T) + + @test isapprox( + dh_dT_p_analytical, + dh_dT_p_AD; + rtol = FT(5e-2), + atol = FT(1e-6), + ) + + # Test ∂θ_li_∂T_sat_p + dθ_li_dT_p_analytical = TD.∂θ_li_∂T_sat_p(param_set, T, p, q_tot) + + f_θ_li_sat_p(T_) = begin + ρ_ = TD.air_density(param_set, T_, p, q_tot) + (q_liq_, q_ice_) = TD.condensate_partition(param_set, T_, ρ_, q_tot) + TD.liquid_ice_pottemp_given_pressure( + param_set, + T_, + p, + q_tot, + q_liq_, + q_ice_, + ) + end + dθ_li_dT_p_AD = ForwardDiff.derivative(f_θ_li_sat_p, T) + + @test isapprox( + dθ_li_dT_p_analytical, + dθ_li_dT_p_AD; + rtol = FT(5e-2), + atol = FT(1e-6), + ) + end + end + + # Test derivatives at fixed density + for T in T_range, ρ in ρ_range, q_tot in q_tot_vals + @testset "∂θ_li_∂T_sat_ρ and ∂p_∂T_sat_ρ: T=$T, ρ=$ρ, q_tot=$q_tot" begin + # Test ∂θ_li_∂T_sat_ρ + dθ_li_dT_ρ_analytical = TD.∂θ_li_∂T_sat_ρ(param_set, T, ρ, q_tot) + + f_θ_li_sat_ρ(T_) = begin + (q_liq_, q_ice_) = TD.condensate_partition(param_set, T_, ρ, q_tot) + TD.liquid_ice_pottemp(param_set, T_, ρ, q_tot, q_liq_, q_ice_) + end + dθ_li_dT_ρ_AD = ForwardDiff.derivative(f_θ_li_sat_ρ, T) + + @test isapprox( + dθ_li_dT_ρ_analytical, + dθ_li_dT_ρ_AD; + rtol = FT(5e-2), + atol = FT(1e-6), + ) + + # Test ∂p_∂T_sat_ρ + dp_dT_ρ_analytical = TD.∂p_∂T_sat_ρ(param_set, T, ρ, q_tot) + + f_p_sat_ρ(T_) = begin + (q_liq_, q_ice_) = TD.condensate_partition(param_set, T_, ρ, q_tot) + TD.air_pressure(param_set, T_, ρ, q_tot, q_liq_, q_ice_) + end + dp_dT_ρ_AD = ForwardDiff.derivative(f_p_sat_ρ, T) + + @test isapprox( + dp_dT_ρ_analytical, + dp_dT_ρ_AD; + rtol = FT(5e-2), + atol = FT(1e-6), + ) + end + end end @testset "∂q_vap_sat_∂T method consistency" begin diff --git a/test/default_saturation_adjustment.jl b/test/default_saturation_adjustment.jl index cfb34b1e..5a515928 100644 --- a/test/default_saturation_adjustment.jl +++ b/test/default_saturation_adjustment.jl @@ -60,7 +60,6 @@ explicit solver method type) work correctly across all formulations: inp.q0, ) - @test res.converged err = abs(res.T - inp.T0) max_err = max(max_err, err) push!(chunk_errs, err) @@ -135,7 +134,7 @@ explicit solver method type) work correctly across all formulations: # Test pe default method for i in idxs inp = targets(i) - let (; T, q_liq, q_ice, converged) = TD.saturation_adjustment( + let (; T, q_liq, q_ice) = TD.saturation_adjustment( param_set, TD.pe(), inp.p0, @@ -143,7 +142,6 @@ explicit solver method type) work correctly across all formulations: inp.q0, ) @test isfinite(T) - @test converged @test isapprox( T, inp.T0; @@ -158,7 +156,7 @@ explicit solver method type) work correctly across all formulations: # Test ph default method for i in idxs inp = targets(i) - let (; T, q_liq, q_ice, converged) = TD.saturation_adjustment( + let (; T, q_liq, q_ice) = TD.saturation_adjustment( param_set, TD.ph(), inp.p0, @@ -166,7 +164,6 @@ explicit solver method type) work correctly across all formulations: inp.q0, ) @test isfinite(T) - @test converged @test isapprox( T, inp.T0; @@ -181,7 +178,7 @@ explicit solver method type) work correctly across all formulations: # Test pρ default method for i in idxs inp = targets(i) - let (; T, q_liq, q_ice, converged) = TD.saturation_adjustment( + let (; T, q_liq, q_ice) = TD.saturation_adjustment( param_set, TD.pρ(), inp.p_ρ, @@ -189,7 +186,6 @@ explicit solver method type) work correctly across all formulations: inp.q0, ) @test isfinite(T) - @test converged @test isapprox( T, inp.T0; @@ -203,7 +199,7 @@ explicit solver method type) work correctly across all formulations: # Test pθ_li default method for i in idxs inp = targets(i) - let (; T, q_liq, q_ice, converged) = TD.saturation_adjustment( + let (; T, q_liq, q_ice) = TD.saturation_adjustment( param_set, TD.pθ_li(), inp.p0, @@ -211,7 +207,6 @@ explicit solver method type) work correctly across all formulations: inp.q0, ) @test isfinite(T) - @test converged @test isapprox( T, inp.T0; @@ -226,7 +221,7 @@ explicit solver method type) work correctly across all formulations: # Test ρθ_li default method for i in idxs inp = targets(i) - let (; T, q_liq, q_ice, converged) = TD.saturation_adjustment( + let (; T, q_liq, q_ice) = TD.saturation_adjustment( param_set, TD.ρθ_li(), inp.ρ0, @@ -234,7 +229,6 @@ explicit solver method type) work correctly across all formulations: inp.q0, ) @test isfinite(T) - @test converged @test isapprox( T, inp.T0; diff --git a/test/runtests_gpu.jl b/test/runtests_gpu.jl index 1512498b..8ab76cb8 100644 --- a/test/runtests_gpu.jl +++ b/test/runtests_gpu.jl @@ -157,6 +157,24 @@ end ql_pe = map(x -> x.q_liq, gpu_results_pe) qi_pe = map(x -> x.q_ice, gpu_results_pe) + # 2b) pe - forced_fixed_iters + gpu_results_pe_fixed = + TD.saturation_adjustment.( + RS.NewtonsMethod, + Ref(param_set), + Ref(TD.pe()), + d_p, + d_e_int_p, + d_q, + Ref(3), + Ref(tol), + nothing, + true, + ) + T_pe_fixed = map(x -> x.T, gpu_results_pe_fixed) + ql_pe_fixed = map(x -> x.q_liq, gpu_results_pe_fixed) + qi_pe_fixed = map(x -> x.q_ice, gpu_results_pe_fixed) + # 3) ph - using SecantMethod (via type dispatch for GPU compatibility) gpu_results_ph = TD.saturation_adjustment.( @@ -173,6 +191,24 @@ end ql_ph = map(x -> x.q_liq, gpu_results_ph) qi_ph = map(x -> x.q_ice, gpu_results_ph) + # 3b) ph - forced_fixed_iters + gpu_results_ph_fixed = + TD.saturation_adjustment.( + RS.NewtonsMethod, + Ref(param_set), + Ref(TD.ph()), + d_p, + d_h_p, + d_q, + Ref(3), + Ref(tol), + nothing, + true, + ) + T_ph_fixed = map(x -> x.T, gpu_results_ph_fixed) + ql_ph_fixed = map(x -> x.q_liq, gpu_results_ph_fixed) + qi_ph_fixed = map(x -> x.q_ice, gpu_results_ph_fixed) + # 4) pρ - using SecantMethod (via type dispatch for GPU compatibility) gpu_results_pρ = TD.saturation_adjustment.( @@ -189,6 +225,24 @@ end ql_pρ = map(x -> x.q_liq, gpu_results_pρ) qi_pρ = map(x -> x.q_ice, gpu_results_pρ) + # 4b) pρ - forced_fixed_iters + gpu_results_pρ_fixed = + TD.saturation_adjustment.( + RS.NewtonsMethod, + Ref(param_set), + Ref(TD.pρ()), + d_p_ρ, + d_ρ, + d_q, + Ref(3), + Ref(tol), + nothing, + true, + ) + T_pρ_fixed = map(x -> x.T, gpu_results_pρ_fixed) + ql_pρ_fixed = map(x -> x.q_liq, gpu_results_pρ_fixed) + qi_pρ_fixed = map(x -> x.q_ice, gpu_results_pρ_fixed) + # 5) ρθ_li - using SecantMethod (via type dispatch for GPU compatibility) gpu_results_ρθ_li = TD.saturation_adjustment.( @@ -205,6 +259,24 @@ end ql_ρθ_li = map(x -> x.q_liq, gpu_results_ρθ_li) qi_ρθ_li = map(x -> x.q_ice, gpu_results_ρθ_li) + # 5b) ρθ_li - forced_fixed_iters + gpu_results_ρθ_li_fixed = + TD.saturation_adjustment.( + RS.NewtonsMethod, + Ref(param_set), + Ref(TD.ρθ_li()), + d_ρ, + d_θ_ρ, + d_q, + Ref(3), + Ref(tol), + nothing, + true, + ) + T_ρθ_li_fixed = map(x -> x.T, gpu_results_ρθ_li_fixed) + ql_ρθ_li_fixed = map(x -> x.q_liq, gpu_results_ρθ_li_fixed) + qi_ρθ_li_fixed = map(x -> x.q_ice, gpu_results_ρθ_li_fixed) + # 6) pθ_li - using SecantMethod (via type dispatch for GPU compatibility) gpu_results_pθ_li = TD.saturation_adjustment.( @@ -221,6 +293,24 @@ end ql_pθ_li = map(x -> x.q_liq, gpu_results_pθ_li) qi_pθ_li = map(x -> x.q_ice, gpu_results_pθ_li) + # 6b) pθ_li - forced_fixed_iters + gpu_results_pθ_li_fixed = + TD.saturation_adjustment.( + RS.NewtonsMethod, + Ref(param_set), + Ref(TD.pθ_li()), + d_p, + d_θ_p, + d_q, + Ref(3), + Ref(tol), + nothing, + true, + ) + T_pθ_li_fixed = map(x -> x.T, gpu_results_pθ_li_fixed) + ql_pθ_li_fixed = map(x -> x.q_liq, gpu_results_pθ_li_fixed) + qi_pθ_li_fixed = map(x -> x.q_ice, gpu_results_pθ_li_fixed) + # Convert GPU results to CPU arrays for comparison T_ρe_cpu = Array(T_ρe) ql_ρe_cpu = Array(ql_ρe) @@ -228,6 +318,21 @@ end T_ρe_fixed_cpu = Array(T_ρe_fixed) ql_ρe_fixed_cpu = Array(ql_ρe_fixed) qi_ρe_fixed_cpu = Array(qi_ρe_fixed) + T_pe_fixed_cpu = Array(T_pe_fixed) + ql_pe_fixed_cpu = Array(ql_pe_fixed) + qi_pe_fixed_cpu = Array(qi_pe_fixed) + T_ph_fixed_cpu = Array(T_ph_fixed) + ql_ph_fixed_cpu = Array(ql_ph_fixed) + qi_ph_fixed_cpu = Array(qi_ph_fixed) + T_pρ_fixed_cpu = Array(T_pρ_fixed) + ql_pρ_fixed_cpu = Array(ql_pρ_fixed) + qi_pρ_fixed_cpu = Array(qi_pρ_fixed) + T_ρθ_li_fixed_cpu = Array(T_ρθ_li_fixed) + ql_ρθ_li_fixed_cpu = Array(ql_ρθ_li_fixed) + qi_ρθ_li_fixed_cpu = Array(qi_ρθ_li_fixed) + T_pθ_li_fixed_cpu = Array(T_pθ_li_fixed) + ql_pθ_li_fixed_cpu = Array(ql_pθ_li_fixed) + qi_pθ_li_fixed_cpu = Array(qi_pθ_li_fixed) T_pe_cpu = Array(T_pe) ql_pe_cpu = Array(ql_pe) qi_pe_cpu = Array(qi_pe) @@ -311,6 +416,11 @@ end end end + # Verify pe fixed iter results + @test isapprox(T_pe_fixed_cpu[i], T_ref; atol = FT(0.05)) + @test isapprox(ql_pe_fixed_cpu[i], ql_ref; atol = FT(1e-4)) + @test isapprox(qi_pe_fixed_cpu[i], qi_ref; atol = FT(1e-4)) + res = TD.saturation_adjustment( RS.SecantMethod, param_set, @@ -341,6 +451,11 @@ end end end + # Verify ph fixed iter results + @test isapprox(T_ph_fixed_cpu[i], T_ref; atol = FT(0.05)) + @test isapprox(ql_ph_fixed_cpu[i], ql_ref; atol = FT(1e-4)) + @test isapprox(qi_ph_fixed_cpu[i], qi_ref; atol = FT(1e-4)) + res = TD.saturation_adjustment( RS.SecantMethod, param_set, @@ -375,6 +490,11 @@ end @test approx_tight(qi_pρ_cpu[i], qi_chk) end + # Verify pρ fixed iter results + @test isapprox(T_pρ_fixed_cpu[i], T_ref; atol = FT(0.05)) + @test isapprox(ql_pρ_fixed_cpu[i], ql_ref; atol = FT(1e-4)) + @test isapprox(qi_pρ_fixed_cpu[i], qi_ref; atol = FT(1e-4)) + res = TD.saturation_adjustment( RS.SecantMethod, param_set, @@ -410,6 +530,11 @@ end @test approx_tight(qi_ρθ_li_cpu[i], qi_chk) end + # Verify ρθ_li fixed iter results + @test isapprox(T_ρθ_li_fixed_cpu[i], T_ref; atol = FT(0.05)) + @test isapprox(ql_ρθ_li_fixed_cpu[i], ql_ref; atol = FT(1e-4)) + @test isapprox(qi_ρθ_li_fixed_cpu[i], qi_ref; atol = FT(1e-4)) + res = TD.saturation_adjustment( RS.SecantMethod, param_set, @@ -447,6 +572,11 @@ end end end + # Verify pθ_li fixed iter results + @test isapprox(T_pθ_li_fixed_cpu[i], T_ref; atol = FT(0.05)) + @test isapprox(ql_pθ_li_fixed_cpu[i], ql_ref; atol = FT(1e-4)) + @test isapprox(qi_pθ_li_fixed_cpu[i], qi_ref; atol = FT(1e-4)) + # Basic invariants for all formulations @test ql_ρe_cpu[i] >= 0 @test qi_ρe_cpu[i] >= 0