Skip to content

Reactant primitive models #985

Draft
maximilian-gelbrecht wants to merge 27 commits intomg/reactantfrom
mg/reactant-primitive
Draft

Reactant primitive models #985
maximilian-gelbrecht wants to merge 27 commits intomg/reactantfrom
mg/reactant-primitive

Conversation

@maximilian-gelbrecht
Copy link
Member

I don't want one monster PR, so I split the Reactant work a little.

#970 is about all the basics and getting the barotropic model going

Here, I continue with the primitive wet model

@maximilian-gelbrecht maximilian-gelbrecht added gpu 🖼️ Everthing GPU related array types 🔢 LowerTriangularMatrices and RingGrids differentiability 🤖 Making the model differentiable via AD labels Feb 9, 2026
@maximilian-gelbrecht
Copy link
Member Author

maximilian-gelbrecht commented Feb 18, 2026

One thing that's annoying to deal with for Reactant compat is the initialize! of components that in the initialize construct some helper struct and then compute something with it. That's the case for e.g. orography with the interpolation, but also the land bucket model that first constructs a seasonal climatology land model.

I don't have the best way to deal with that figured out yet without going directly into the initialize! of these components and selectively apply a @maybe_jit

Comment on lines +13 to +17
# Branchless kernel for finding ring indices, compatible with Reactant tracing.
# It's slower than the regular kernel on both CPU and GPU, that's why it's only used here
# Counts how many latd entries are >= θ to find the ring index j
# such that latd[j] >= θ > latd[j+1]. This avoids data-dependent control flow
# (while, break, if/elseif) which is incompatible with Reactant tracing.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is good to know for the parameterizations. For a lot of the surface fluxes we use condition ? a : b which can be easily written as ifelse but I'll try to directly implement this also for the convection scheme which has way more conditionals that I'll try to revisit in #976

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Branches are not completely forbidding, we can @trace them in Reactant as well. But the binary search algorithm here was problematic with its data dependent while loop.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

array types 🔢 LowerTriangularMatrices and RingGrids differentiability 🤖 Making the model differentiable via AD gpu 🖼️ Everthing GPU related

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants