Fix moisture partitioning and unify moist-state conversions - #997
kaiyuan-cheng wants to merge 14 commits into
Conversation
Convert total moisture `qᵗ` to the scheme-dependent prognostic moisture `qᵛᵉ` with one
generic method per argument form, driven by `condensate_field_names`. Those names are the
complement of the prognostic moisture — the partition `total_condensate_density` already
uses — so the same expression recovers vapor for non-equilibrium schemes and equilibrium
moisture `qᵉ` for saturation adjustment, with no per-scheme method.
This replaces a P3-specific `μ::NamedTuple` override and seven hand-written state-based
overloads (Kessler, P3, WP1M, MP1M, WPNE1M, MPNE1M, WPNE2M), which were a second source of
truth and had already drifted: five clamped the result at zero and two did not. The
conversion now preserves total water without clipping, so a bad initial condition surfaces
instead of silently gaining water; `set!` validates it against a tolerance relative to the
water in each cell, and the parcel correction keeps condensate within the budget during
time stepping.
`set!(::AtmosphereModel)` drops the `hasmethod` guard that skipped the `qᵗ` conversion
entirely: only P3 had a four-argument method, so every other scheme left `qᵗ` sitting in
the `qᵛᵉ` slot while the microphysical fields held the same water again, double-counting
condensate in the initial state. The conversion and its validation move to
`set_moisture.jl`; under Reactant the validating reduction is a traced Boolean, so it is
checked in a runtime callback rather than while tracing.
Saturation adjustment now retains prognostic precipitation. Previously the adjusted state
was rebuilt from `qᵉ` alone, so rain and snow vanished from the partition: the diagnosed
cloud liquid was cloud minus rain, and every autoconversion or accretion event appeared as
a spurious cooling of `ℒΔq/cᵖ` because `qᵉ` fell while `θˡⁱ` stayed fixed. Precipitation
now enters the heat capacity, gas constant, and latent terms, which is the convention the
non-equilibrium schemes already followed. The pressure-based and density-based adjustments
collapse into one secant iteration, with the state selecting the saturation constraint and
the residual, removing the duplicate `LiquidIceDensityState` implementation.
Fix two moist-Exner inconsistencies the generalized helpers expose:
- `wall_potential_temperature` computed the wall `θ₀` with the dry `Rᵈ/cᵖᵈ` while
differencing it against a moist `θˡⁱ`, biasing the bulk sensible heat flux. It now
uses the near-wall composition, as `wall_static_energy` already did.
- The parcel `θ → T` kernel was also dry. `set!(::ParcelModel)` now sets moisture before
temperature, and the relative-humidity path, which needs `T` to diagnose saturation,
takes a preliminary pass and repeats the conversion.
`temperature_from_potential_temperature` and its inverse accept full `MoistureMassFractions`
so condensate enters both the Exner function and the latent-heat term.
Surface filtering: `filter_timescale = Inf` is the default, and it drove `ϵ = Δt/τ` to
zero, freezing the filtered fields at whatever `initialize!` wrote, so a bulk flux could
not see the wind or the wall temperature change. It now samples the current state, meaning
"no temporal averaging" rather than "no update at all". Both bulk scalar fluxes store the
complete near-wall difference, wall value included, so the wall and the air are sampled at
the same times; filtering only the atmospheric half would have pinned `T₀` at its initial
value, and left the heat and vapor fluxes disagreeing about when the wall was observed.
The parcel state is initialized consistently with the condensate its prognostics carry:
the water is partitioned before the static energy is formed, so a parcel holding rain
starts at the environmental temperature with the right heat capacity and latent term
instead of waiting for the first substep to rewrite it. `fix_negative_moisture!` gains a
parcel method that keeps condensate nonnegative and within the conserved total water.
`specific_field_name` builds strings, which allocates and does not constant-fold, so the
density-to-specific name mapping is generated at compile time rather than derived inside
the kernels that use it.
Bulk boundary conditions are materialized before the dynamics, so the sensible heat flux
captures whatever `total_density` returned on the dynamics stub. That holds today, but
nothing enforced it, and a dynamics that allocated its density during materialization
would hand the boundary condition a stale field with no error. Check it at `initialize!`
and record the invariant in the `materialize_dynamics` docstring.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Note
Copilot was unable to run its full agentic suite in this review.
Pull request overview
This PR improves moisture-aware thermodynamics and boundary-condition handling, with expanded test coverage for total-water initialization, wall flux filtering behavior, and saturation adjustment retaining precipitation.
Changes:
- Make potential temperature / temperature conversions use mixture thermodynamics (vapor + condensates) via
MoistureMassFractions. - Refactor total-water (
qᵗ) initialization to convert to scheme-specific prognostic moisture generically, validating against condensate budgets (including Reactant execution-time validation). - Update bulk scalar flux filtering to filter the complete wall–air difference and adjust saturation adjustment to retain prognostic precipitation during adjustment.
Reviewed changes
Copilot reviewed 29 out of 29 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| test/wall_fluxes.jl | Adds wall-flux tests for moisture-fixed sensible heat, sampling behavior, and vapor filtering. |
| test/reactant/moisture_initialization.jl | Adds Reactant + Enzyme tests ensuring compiled initialization reads new inputs and remains differentiable. |
| test/polynomial_bulk_coefficients.jl | Adds regression tests for filter_timescale=Inf direct-sampling (no averaging) behavior. |
| test/moist_state_conversion.jl | Adds extensive tests for moist thermodynamic conversions and total-water conversions across schemes. |
| test/forcing_and_boundary_conditions.jl | Adjusts timestep used in boundary-condition forcing test. |
| src/Thermodynamics/thermodynamics_constants.jl | Adds typed conversion constructor for MoistureMassFractions{FT}. |
| src/Thermodynamics/dynamic_states.jl | Updates θ↔T convenience functions to use mixture properties via MoistureMassFractions. |
| src/StaticEnergyFormulations/static_energy_tendency.jl | Passes microphysical prognostics into thermodynamic adjustment path. |
| src/PotentialTemperatureFormulations/potential_temperature_tendency.jl | Passes microphysical prognostics into thermodynamic adjustment path. |
| src/ParcelModels/parcel_dynamics.jl | Improves parcel initialization to preserve θ & RH, partitions moisture consistently, and retains precipitation during adjustment. |
| src/Microphysics/saturation_adjustment.jl | Unifies saturation adjustment across pressure/density states and supports fixed precipitation during adjustment. |
| src/Microphysics/dcmip2016_kessler.jl | Removes scheme-specific total→prognostic moisture conversion in favor of generic interface. |
| src/Microphysics/PredictedParticleProperties/p3_microphysical_state.jl | Removes P3-specific total→prognostic moisture conversion in favor of generic interface. |
| src/BoundaryConditions/update_boundary_conditions.jl | Changes filtered scalar “source” to a kernel operation returning the full wall difference. |
| src/BoundaryConditions/thermodynamic_variable_bcs.jl | Ensures sensible-heat formulation updates preserve new moisture capture. |
| src/BoundaryConditions/filtered_surface_state.jl | Refactors filtering to support filter_timescale=Inf as direct sampling; updates docs. |
| src/BoundaryConditions/bulk_scalar_fluxes.jl | Makes wall thermodynamic values moisture-aware and changes filtered scalar semantics to store complete differences. |
| src/BoundaryConditions/BoundaryConditions.jl | Captures moisture context (microphysics + density handle) when materializing sensible heat BCs. |
| src/AtmosphereModels/update_atmosphere_model_state.jl | Passes prognostic precipitation info into thermodynamic adjustment. |
| src/AtmosphereModels/set_moisture.jl | New helper to convert staged total-water density to scheme prognostic moisture with validation. |
| src/AtmosphereModels/set_atmosphere_model.jl | Routes total-water inputs through convert_total_moisture! and removes per-scheme conversion plumbing. |
| src/AtmosphereModels/microphysics_interface.jl | Adds compile-time name resolution and generic total→prognostic moisture conversion utilities. |
| src/AtmosphereModels/dynamics_interface.jl | Documents BC materialization ordering constraints around density fields. |
| src/AtmosphereModels/AtmosphereModels.jl | Includes new set_moisture.jl. |
| ext/BreezeReactantExt/initialization.jl | Adds Reactant runtime callback specialization for moisture validation. |
| ext/BreezeReactantExt/BreezeReactantExt.jl | Includes new Reactant initialization extension file. |
| ext/BreezeCloudMicrophysicsExt/two_moment_microphysics.jl | Removes scheme-specific total→prognostic moisture conversion method. |
| ext/BreezeCloudMicrophysicsExt/one_moment_microphysics.jl | Updates saturation-adjustment hook to retain precipitation during thermodynamic adjustment. |
| docs/src/developer/microphysics/overview.md | Documents new total-water conversion and extended thermodynamic adjustment interface. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| function sensible_heat_source_field(bf, model) | ||
| return KernelFunctionOperation{Center, Center, Center}(sensible_heat_difference, model.grid, | ||
| bf, model.clock, Oceananigans.fields(model)) | ||
| end | ||
|
|
||
| @inline function sensible_heat_difference(i, j, k, grid, bf, clock, fields) | ||
| T₀ = wall_value(i, j, grid, Bottom(), bf.surface_temperature, clock) | ||
| return bulk_sensible_heat_difference(i, j, k, grid, Bottom(), bf.formulation, bf, T₀, fields, nothing) | ||
| end | ||
|
|
||
| vapor_source_field(model) = AtmosphereModels.specific_prognostic_moisture(model) | ||
| function vapor_source_field(bf, model) | ||
| return KernelFunctionOperation{Center, Center, Center}(vapor_difference, model.grid, | ||
| bf, model.clock, Oceananigans.fields(model)) | ||
| end | ||
|
|
||
| @inline function vapor_difference(i, j, k, grid, bf, clock, fields) | ||
| T₀ = wall_value(i, j, grid, Bottom(), bf.surface_temperature, clock) | ||
| qᵛ₀ = wall_specific_humidity(i, j, grid, Bottom(), bf, T₀, clock) | ||
| return bulk_vapor_difference(i, j, k, fields, nothing, qᵛ₀) | ||
| end |
| function validate_total_moisture(total_moisture, prognostic_moisture, validation_field) | ||
| margin = KernelFunctionOperation{Center, Center, Center}(total_moisture_validation_margin, | ||
| validation_field.grid, total_moisture, prognostic_moisture) | ||
| # Reactant reductions need a stored field. Reuse the specific moisture field | ||
| # as scratch; the conversion above restores it before storing the density. | ||
| set!(validation_field, margin) | ||
| validate_total_moisture(minimum(validation_field) ≥ 0) | ||
| return nothing | ||
| end | ||
|
|
||
| # Traced backends specialize this scalar check to run at execution time. | ||
| function validate_total_moisture(valid) | ||
| Bool(valid) || throw(ArgumentError("set! received a total moisture qᵗ smaller than the supplied \ | ||
| condensates, or a non-finite moisture state. Increase qᵗ or \ | ||
| reduce the condensate inputs.")) | ||
| return nothing | ||
| end |
| throw(ArgumentError("The density captured by a BulkSensibleHeatFlux boundary condition is not \ | ||
| the model's total density. Boundary conditions are materialized before \ | ||
| the dynamics, so `total_density` of the dynamics stub must be either \ | ||
| `nothing` or the same field the materialized dynamics carries.")) |
condensate_field_namesPreserve moist-state conversion, precipitation retention, moist sensible heat exchange, and complete temporal filtering while integrating main's live wall-pressure interface. Adapt the snow prognostic to qˢⁿ throughout the new saturation-adjustment path and its conservation tests. Keep both parents' regression coverage and update merged wall-flux expectations for live pressure and the model clock's precision.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
| q₁ = MoistureMassFractions(qᵉ) | ||
| @inline function AM.maybe_adjust_thermodynamic_state(𝒰₀, bμp::Union{WP1M, MP1M}, qᵉ, constants, μ, ρ) | ||
| qʳ = μ.ρqʳ / ρ | ||
| qˢⁿ = get(μ, :ρqˢⁿ, zero(ρ)) / ρ |
There was a problem hiding this comment.
can we use dispatch here, not get ?
There was a problem hiding this comment.
I just realized that I missed a few of your comments. This is done.
| # Resolve names to literals at compile time: string processing cannot run inside GPU kernels. | ||
| # QuoteNode keeps a scalar Symbol literal from being interpreted as a variable name. | ||
| @generated specific_field_name(::Val{name}) where name = QuoteNode(specific_field_name(name)) | ||
| @generated specific_field_names(::Val{names}) where names = :($(map(specific_field_name, names))) |
There was a problem hiding this comment.
My understand is that these allow us to use the specific names inside kernels. moisture_specific_name and specific_condensate_names are now called from kernel code (e.g., , fields[moisture_specific_name(microphysics)][i, j, k] in wall_moisture_fractions). The Symbol method strips the ρ using string operations, which can't run in a GPU kernel. The @generated Val methods do that stripping at compile time, so the kernel only ever sees a literal name. QuoteNode keeps that literal Symbol from being read as a variable.
| @inline moisture_less_condensate(qᵗ, ℳ, names::Tuple{Symbol, Vararg}) = | ||
| qᵗ - sum_microphysical_components(ℳ, names) |
There was a problem hiding this comment.
| @inline moisture_less_condensate(qᵗ, ℳ, names::Tuple{Symbol, Vararg}) = | |
| qᵗ - sum_microphysical_components(ℳ, names) | |
| @inline moisture_less_condensate(qᵗ, ℳ, names::Tuple{Symbol, Vararg}) = qᵗ - sum_microphysical_components(ℳ, names) |
There was a problem hiding this comment.
I don't totally understand the names. "moisture less condensate" is q^t - q^c ? is that different from the vapor mass fraction?
There was a problem hiding this comment.
This helper subtracts the species listed by the microphysics scheme:
- Non-equilibrium microphysics: subtract all hydrometeors. The result is vapor mass fraction.
- Saturation adjustment: subtract only precipitating hydrometeors. The result is equilibrium moisture, qᵉ = qᵛ + qᶜˡ + qᶜⁱ, which still includes cloud droplet and cloud ice.
We can come up with a better name.
There was a problem hiding this comment.
I will go with subtract_condensate. Feel free to suggest a new name.
|
|
||
| microphysics = DCMIP2016KesslerMicrophysics() | ||
| μ = (ρqᶜˡ=0.0012, ρqʳ=0.0024) | ||
| specific_prognostic_moisture_from_total(microphysics, 0.02, μ, 1.2) |
There was a problem hiding this comment.
why do we need the suffix _from_total here? Is there another competing method specific_prognostic_moisture that we need to distinguish?
There was a problem hiding this comment.
Yes. There is specific_prognostic_moisture(model), which retrieves the model’s existing qᵛ or qᵉ field.
specific_prognostic_moisture_from_total instead converts supplied total water into qᵛ or qᵉ by subtracting the appropriate condensates.
We could use dispatch so that we don't need the suffix.
|
|
||
| # Sum scalars, whole fields, or state components. Stop at the last component to avoid | ||
| # adding a scalar zero to lazy field expressions; moisture_less_condensate handles empty tuples. | ||
| @inline sum_microphysical_components(μ, names::Tuple{Symbol}) = getproperty(μ, first(names)) |
There was a problem hiding this comment.
this method doesn't seem specific to microphysics
There was a problem hiding this comment.
You are right. Moved this into Breeze.Utils.
Co-authored-by: Gregory L. Wagner <gregory.leclaire.wagner@gmail.com>
Co-authored-by: Gregory L. Wagner <gregory.leclaire.wagner@gmail.com>
Co-authored-by: Gregory L. Wagner <gregory.leclaire.wagner@gmail.com>
|
The 7 CI jobs that failed on 8313f0d all trace back to two of the new tests, not to the source changes. The same test code is still on 68849ee, so the run now in progress should fail the same way. With the patch below,
--- a/test/wall_fluxes.jl
+++ b/test/wall_fluxes.jl
@@ -488,7 +488,7 @@
- T_wall(t) = FT(290) + FT(10) * sinpi(t / FT(20))
+ T_wall(t) = 290 + 10 * sinpi(t / 20)
@@ -562,7 +562,7 @@
- T_wall(t) = FT(290) + FT(10) * sinpi(t / FT(20))
+ T_wall(t) = 290 + 10 * sinpi(t / 20)
--- a/test/moist_state_conversion.jl
+++ b/test/moist_state_conversion.jl
@@ -66 @@ Parcel initialization preserves θ and relative humidity
- grid = RectilinearGrid(default_arch, FT; size=4, z=(0, 100), topology=(Flat, Flat, Bounded))
+ grid = RectilinearGrid(CPU(), FT; size=4, z=(0, 100), topology=(Flat, Flat, Bounded))
@@ -140 @@ Parcel condensate overshoots preserve water and energy
- grid = RectilinearGrid(default_arch, FT; size=4, z=(0, 100), topology=(Flat, Flat, Bounded))
+ grid = RectilinearGrid(CPU(), FT; size=4, z=(0, 100), topology=(Flat, Flat, Bounded))
@@ -256 @@ Parcel initialization partitions carried condensate
- grid = RectilinearGrid(default_arch, FT; size=4, z=(0, 100), topology=(Flat, Flat, Bounded))
+ grid = RectilinearGrid(CPU(), FT; size=4, z=(0, 100), topology=(Flat, Flat, Bounded))The Julia 1.13 logs also show |
It was Reactant, but in the CPU tests we dont need Reactant and so the error was non-fatal. Also, that's been resolved by #923, so I presume this refers to an old run. |
There was a problem hiding this comment.
⚠️ Performance Alert ⚠️
Possible performance regression was detected for benchmark 'Breeze.jl Benchmarks'.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 1.10.
| Benchmark suite | Current: 68849ee | Previous: d0afaf4 | Ratio |
|---|---|---|---|
ScalarTendency; Grid: 256x256x128/Advection: WENO5/NVIDIA L4/F32 vanilla |
6745119568.087938 points/s |
7633485225.69149 points/s |
1.13 |
ScalarTendency; Grid: 256x256x128/Advection: WENO5/NVIDIA L4/F32 reactant raise=true |
7742811097.594069 points/s |
8613789690.07802 points/s |
1.11 |
This comment was automatically generated by workflow using github-action-benchmark.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Closes #1013
Summary
Total moisture
qᵗis converted to a scheme's prognostic moistureqᵛᵉby one genericmethod per argument form, driven by
condensate_field_names. Those names are thecomplement of the prognostic moisture, the same partition
total_condensate_densityalready uses, so a single expression recovers vapor for non-equilibrium schemes and
equilibrium moisture
qᵉfor saturation adjustment, with no per-scheme method.This replaces a P3-specific
μ::NamedTupleoverride and seven hand-written state-basedoverloads (Kessler, P3, WP1M, MP1M, WPNE1M, MPNE1M, WPNE2M). They were a second source of
truth and had already drifted: five clamped the result at zero and two did not.
Results change
This is a correctness fix rather than a refactor, and existing runs will produce different
numbers:
temperature and the diagnosed cloud liquid, and removes the spurious cooling that
accompanied every autoconversion and accretion event.
is moist, and both bulk scalar fluxes now sample the wall at the same time as the air.
filter_timescale = Inf. The filtered surface statewas frozen after initialization and now follows the current state, so bulk fluxes
respond to the wind and the wall temperature.
Initial conditions also change wherever
qᵗwas supplied together with condensate, sincethat water was previously counted twice.
Bugs fixed
set!guarded theqᵗconversion with
hasmethod, and only P3 had a four-argument method. Every other schemeleft
qᵗin theqᵛᵉslot while the microphysical fields held the same water again.qᵉalone, so the diagnosed cloud liquid was cloud minus rain, and every autoconversionor accretion event appeared as a spurious cooling of
ℒΔq/cᵖbecauseqᵉfell whileθˡⁱstayed fixed. Precipitation now enters the heat capacity, gas constant, and latentterms, the convention the non-equilibrium schemes already followed.
wall_potential_temperatureformed
θ₀withRᵈ/cᵖᵈand differenced it against a moistθˡⁱ. It now uses thenear-wall composition, as
wall_static_energyalready did.θ → Tconversion was dry.set!(::ParcelModel)now sets moisture first;the relative-humidity path, which needs
Tto diagnose saturation, takes a preliminarypass and repeats the conversion.
filter_timescale = Inf, the default, froze the filtered surface fields. It droveϵ = Δt/τto zero, so a bulk flux never saw the wind or the wall temperature changeafter initialization. It now samples the current state, meaning "no temporal averaging"
rather than "no update at all".
complete near-wall difference, wall value included, so with a time-varying surface
temperature the heat and vapor fluxes agree about when the wall was sampled.
before the static energy is formed, so a parcel holding rain starts at the environmental
temperature with the right heat capacity and latent term instead of waiting for the first
substep to rewrite it.
Other changes
iteration, the state selecting the saturation constraint and the residual. This removes
the duplicate
LiquidIceDensityStateimplementation added in Compressible θˡⁱ thermodynamic state uses a pressure coordinate — one root cause behind both the temperature-inversion inconsistency (#752 / #625-6) and a saturation adjustment evaluated at the wrong pressure #765.surfaces rather than silently gaining water.
set!validates it against a tolerancerelative to the water in each cell; under Reactant the validating reduction is a traced
Boolean, so it is checked in a runtime callback rather than while tracing.
fix_negative_moisture!gains a parcel method keeping condensate nonnegative and withinthe conserved total water.
specific_field_namebuilds strings, which allocates and does not constant-fold, so thedensity-to-specific name mapping is generated at compile time rather than derived inside
the kernels that use it.
CompressibleDynamicserrors:ef.density === nothing(BC materialized beforematerialize_dynamics) #777), so the sensibleheat flux captures whatever
total_densityreturned on the dynamics stub. Nothingenforced that, and a dynamics allocating its density during materialization would hand
the boundary condition a stale field with no error. It is checked at
initialize!andthe invariant is recorded in the
materialize_dynamicsdocstring.temperature_from_potential_temperatureand its inverse accept fullMoistureMassFractions, so condensate enters both the Exner function and the latent term.🤖 Generated with Claude Code