Reland "Test raising one moment microphysics" by giordano · Pull Request #609 · NumericalEarth/Breeze.jl

giordano · 2026-04-03T14:34:34Z

Reverts #608 and relands #595. This will need some work to fix the raising failure caused by #587 + #595, but opening the PR not to forget about it

This reverts commit b7cf547.

codecov · 2026-04-03T14:52:29Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

dkytezab · 2026-04-03T16:39:25Z

@giordano / @kaiyuan-cheng I think i see the issue. #587 added conv_q_icl_to_q_sno_no_supersat from Clima CloudMicrophysics which ends up calling LogExpFunctions.log1pexp. The problem is that we don't currently have a raising rule for the corresponding lowered operation (math.log1p inside of MLIR). So Reactant sees this call inside of a kernel, doesn't know how to deal with it, and then throws an error.

The temporary fix is to set use the non-smooth branch by setting true to false below

Breeze.jl/ext/BreezeCloudMicrophysicsExt/one_moment_microphysics.jl

Lines 844 to 846 in b7cf547

    
           # Ice → snow autoconversion 
        
           Sᵃᶜⁿᵛⁱˢ = conv_q_icl_to_q_sno_no_supersat(categories.snow.acnv1M, qᶜⁱ, true)

cause this avoids the code path. I wonder if we can add the smoothness flag to the mpne1m signature so the user can control this? Long term is getting this rule in Enzyme-JAX which i can look in to

glwagner · 2026-04-03T17:51:35Z

what does "non-smooth branch" mean?

dkytezab · 2026-04-03T18:03:19Z

smoothing the kink

@inline conv_q_icl_to_q_sno_no_supersat(
    (; τ, q_threshold, k)::CMP.Acnv1M{FT},
    q_icl::FT,
    smooth_transition::Bool = false,
) where {FT} =
    smooth_transition ?
    CO.logistic_function_integral(q_icl, q_threshold, k) / τ :
    max(0, q_icl - q_threshold) / τ

glwagner · 2026-04-03T18:08:56Z

smoothing the kink

@inline conv_q_icl_to_q_sno_no_supersat(
    (; τ, q_threshold, k)::CMP.Acnv1M{FT},
    q_icl::FT,
    smooth_transition::Bool = false,
) where {FT} =
    smooth_transition ?
    CO.logistic_function_integral(q_icl, q_threshold, k) / τ :
    max(0, q_icl - q_threshold) / τ

v interesting

…yiscs

giordano · 2026-04-06T22:27:29Z

Still failing with Julia v1.12 only: https://github.com/NumericalEarth/Breeze.jl/actions/runs/24053001251/job/70152934537?pr=609#step:10:632

Error in testset Raise backward:
Error During Test at /__w/Breeze.jl/Breeze.jl/test/reactant_1m_microphysics_compilation.jl:89
  Got exception outside of a @test
  StackOverflowError:

Is there a recursive function somewhere?

…yiscs

giordano · 2026-04-23T21:43:34Z

The stackoverflow error seems to be gone for the time being, but now we have a new exciting error:

Error During Test at /__w/Breeze.jl/Breeze.jl/test/reactant_1m_microphysics_compilation.jl:89
  Got exception outside of a @test
  TypeError: non-boolean (Reactant.TracedRNumber{Bool}) used in boolean context
  Stacktrace:
    [1] ==
      @ ./abstractarray.jl:3037
    [2] call_with_reactant(::Reactant.EnsureReturnType{Any}, ::typeof(==), ::SubArray{Reactant.TracedRNumber{Float64}, 1, Reactant.TracedRArray{Float64, 1}, Tuple{UnitRange{Int64}}, true}, ::SubArray{Reactant.TracedRNumber{Float64}, 1, Reactant.TracedRArray{Float64, 1}, Tuple{UnitRange{Int64}}, true})
      @ Reactant /usr/local/share/julia/packages/Reactant/RR3Dn/src/utils.jl:1190
    [3] macro expansion
      @ /usr/local/share/julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:206 [inlined]
    [4] ==
      @ /usr/local/share/julia/packages/Oceananigans/tBtso/src/Grids/abstract_grid.jl:91
    [5] call_with_reactant(::Reactant.EnsureReturnType{Any}, ::typeof(==), ::Oceananigans.Grids.LatitudeLongitudeGrid{Reactant.TracedRNumber{Float64}, Oceananigans.Grids.Periodic, Oceananigans.Grids.Bounded, Oceananigans.Grids.Bounded,

dkytezab · 2026-04-23T21:47:41Z

Exciting! I've been hitting this elsewhere. Is this validate_grid?

giordano · 2026-04-23T21:49:10Z

    [6] call_with_reactant(::Reactant.EnsureReturnType{Oceananigans.Grids.LatitudeLongitudeGrid{Reactant.TracedRNumber{Float64}, Oceananigans.Grids.Periodic, Oceananigans.Grids.Bounded, Oceananigans.Grids.Bounded, Oceananigans.Grids.StaticVerticalDiscretization{OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, Float64, Float64}, Float64, Float64, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, Float64, Float64, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, Float64, Float64, Oceananigans.Architectures.ReactantState, Int64}}, ::typeof(Oceananigans.AbstractOperations.validate_grid), ::Oceananigans.Fields.Field{Oceananigans.Grids.Center, Oceananigans.Grids.Center, Oceananigans.Grids.Center, Nothing, Oceananigans.Grids.LatitudeLongitudeGrid{Reactant.TracedRNumber{Float64}, Oceananigans.Grids.Periodic, Oceananigans.Grids.Bounded, Oceananigans.Grids.Bounded, Oceananigans.Grids.StaticVerticalDiscretization{OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, Float64, Float64}, Float64, Float64, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, Float64, Float64, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, Float64, Float64, Oceananigans.Architectures.ReactantState, Int64}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Reactant.TracedRNumber{Float64}, 3, Reactant.TracedRArray{Float64, 3}}, Float64, Oceananigans.BoundaryConditions.FieldBoundaryConditions{Oceananigans.BoundaryConditions.BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Oceananigans.BoundaryConditions.BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Nothing, @NamedTuple{bottom_and_top::KernelAbstractions.Kernel{ReactantKernelAbstractionsExt.ReactantBackend, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, Oceananigans.Utils.OffsetStaticSize{(1:8, 1:8)}, typeof(Oceananigans.BoundaryConditions.gpu__fill_bottom_and_top_halo!)}, south_and_north::KernelAbstractions.Kernel{ReactantKernelAbstractionsExt.ReactantBackend, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, Oceananigans.Utils.OffsetStaticSize{(1:8, 1:8)}, typeof(Oceananigans.BoundaryConditions.gpu__fill_south_and_north_halo!)}, west_and_east::Oceananigans.BoundaryConditions.PeriodicFillHalo{KernelAbstractions.Kernel{ReactantKernelAbstractionsExt.ReactantBackend, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, Oceananigans.Utils.OffsetStaticSize{(1:14, 1:14)}, typeof(Oceananigans.BoundaryConditions.gpu__fill_periodic_west_and_east_halo!)}, 8, 3}}, @NamedTuple{bottom_and_top::Tuple{Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition}, south_and_north::Tuple{Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition}, west_and_east::Tuple{Oceananigans.BoundaryConditions.BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Oceananigans.BoundaryConditions.BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}}}}, Nothing, Nothing}, ::Oceananigans.Fields.Field{Oceananigans.Grids.Center, Oceananigans.Grids.Center, Oceananigans.Grids.Center, Nothing, Oceananigans.Grids.LatitudeLongitudeGrid{Reactant.TracedRNumber{Float64}, Oceananigans.Grids.Periodic, Oceananigans.Grids.Bounded, Oceananigans.Grids.Bounded, Oceananigans.Grids.StaticVerticalDiscretization{OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, Float64, Float64}, Float64, Float64, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, Float64, Float64, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, OffsetArrays.OffsetVector{Reactant.TracedRNumber{Float64}, Reactant.TracedRArray{Float64, 1}}, Float64, Float64, Oceananigans.Architectures.ReactantState, Int64}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Reactant.TracedRNumber{Float64}, 3, Reactant.TracedRArray{Float64, 3}}, Float64, Oceananigans.BoundaryConditions.FieldBoundaryConditions{Oceananigans.BoundaryConditions.BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Oceananigans.BoundaryConditions.BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Nothing, @NamedTuple{bottom_and_top::KernelAbstractions.Kernel{ReactantKernelAbstractionsExt.ReactantBackend, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, Oceananigans.Utils.OffsetStaticSize{(1:8, 1:8)}, typeof(Oceananigans.BoundaryConditions.gpu__fill_bottom_and_top_halo!)}, south_and_north::KernelAbstractions.Kernel{ReactantKernelAbstractionsExt.ReactantBackend, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, Oceananigans.Utils.OffsetStaticSize{(1:8, 1:8)}, typeof(Oceananigans.BoundaryConditions.gpu__fill_south_and_north_halo!)}, west_and_east::Oceananigans.BoundaryConditions.PeriodicFillHalo{KernelAbstractions.Kernel{ReactantKernelAbstractionsExt.ReactantBackend, KernelAbstractions.NDIteration.StaticSize{(16, 16)}, Oceananigans.Utils.OffsetStaticSize{(1:14, 1:14)}, typeof(Oceananigans.BoundaryConditions.gpu__fill_periodic_west_and_east_halo!)}, 8, 3}}, @NamedTuple{bottom_and_top::Tuple{Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition}, south_and_north::Tuple{Oceananigans.BoundaryConditions.NoFluxBoundaryCondition, Oceananigans.BoundaryConditions.NoFluxBoundaryCondition}, west_and_east::Tuple{Oceananigans.BoundaryConditions.BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Oceananigans.BoundaryConditions.BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}}}}, Nothing, Nothing})
      @ Reactant /usr/local/share/julia/packages/Reactant/RR3Dn/src/utils.jl:1190
[...]
    [9] #set!#37
      @ /__w/Breeze.jl/Breeze.jl/src/AtmosphereModels/set_atmosphere_model.jl:173

Looks like it

dkytezab · 2026-04-23T23:51:57Z

Breeze.jl/src/AtmosphereModels/set_atmosphere_model.jl

Lines 169 to 173 in 1ccbb30

    
           elseif name ∈ (:ρqᵗ, :ρqᵛ, :ρqᵉ) 
        
               set!(model.moisture_density, value) 
        
               ρ = dynamics_density(model.dynamics) 
        
               qᵛᵉ = specific_prognostic_moisture(model) 
        
               set!(qᵛᵉ, model.moisture_density / ρ)

we could try using parent here instead, or override the grid equality check on LatitudeLongitudeGrid with ReactantState in Oceananigans...

glwagner · 2026-04-24T04:30:04Z

Breeze.jl/src/AtmosphereModels/set_atmosphere_model.jl

Lines 169 to 173 in 1ccbb30

elseif name ∈ (:ρqᵗ, :ρqᵛ, :ρqᵉ)

set!(model.moisture_density, value)

ρ = dynamics_density(model.dynamics)

qᵛᵉ = specific_prognostic_moisture(model)

set!(qᵛᵉ, model.moisture_density / ρ)

we could try using parent here instead, or override the grid equality check on LatitudeLongitudeGrid with ReactantState in Oceananigans...

why do the grids differ?

dkytezab · 2026-04-24T16:44:18Z

They don't, the issue is computing equalities between TracedRNumber

glwagner · 2026-04-24T17:44:19Z

They don't, the issue is computing equalities between TracedRNumber

Do we need a custom equality operator for Reactant then?

Revert "Revert "Test raising one moment microphysics" (#608)"

150bb7a

This reverts commit b7cf547.

giordano assigned dkytezab Apr 3, 2026

dkytezab and others added 6 commits April 3, 2026 12:17

Merge branch 'main' into revert-608-revert-595-dkz/feat-raise-microph…

d04f2d1

…yiscs

Merge branch 'main' into revert-608-revert-595-dkz/feat-raise-microph…

990d1d6

…yiscs

DONT MERGE (test reactant branch with fix)

0687e03

Update Project.toml

4cc37e0

Update Project.toml

4a8aada

Fix projects

f48c9fc

giordano added 2 commits April 10, 2026 13:09

Merge branch 'main' into revert-608-revert-595-dkz/feat-raise-microph…

e81f94b

…yiscs

Merge branch 'main' into revert-608-revert-595-dkz/feat-raise-microph…

935ee65

…yiscs

Conversation

giordano commented Apr 3, 2026

Uh oh!

codecov Bot commented Apr 3, 2026

Codecov Report

Uh oh!

dkytezab commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

glwagner commented Apr 3, 2026

Uh oh!

dkytezab commented Apr 3, 2026

Uh oh!

glwagner commented Apr 3, 2026

Uh oh!

giordano commented Apr 6, 2026

Uh oh!

giordano commented Apr 23, 2026

Uh oh!

dkytezab commented Apr 23, 2026

Uh oh!

giordano commented Apr 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dkytezab commented Apr 23, 2026

Uh oh!

glwagner commented Apr 24, 2026

Uh oh!

dkytezab commented Apr 24, 2026

Uh oh!

glwagner commented Apr 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dkytezab commented Apr 3, 2026 •

edited

Loading

giordano commented Apr 23, 2026 •

edited

Loading