Reactant `LandModel` without vegetation by maximilian-gelbrecht · Pull Request #157 · NumericalEarth/Terrarium.jl

maximilian-gelbrecht · 2026-08-04T07:35:44Z

This ensures Reactant compatibility of a LandModel without vegetation.

bgroenks96 · 2026-08-04T07:53:54Z

 """
 function initialize!(state, grid, ::ImplicitSkinTemperature, args...)
-    set!(state.skin_temperature, state.ground_temperature)
+    interior(state.skin_temperature) .= interior(state.ground_temperature)   # Now set! for better Reactant compat, set! might make a CPU detour


The comment doesn't match the change here. Why do we need to stop using set!? I would rather keep that.

It's just a typo "No set!".

The problem is that set! might make CPU detours and the interpolation isn't fully Reactant compatible.

Did you create an Oceananigans issue about that?

I don't see any reason why set! should make CPU detours when setting a Field to another Field.

No, not yet

Ok, then please do that and link it here, and add a TODO comment in the code.

Indeed, it should have. Except that reactant defines its own set_to_field!

So the somewhat funny situation is that the Reactant extension of set_to_field! actually causes set! to fail with Reactant, when it would otherwise succeed.

It's fixed in a new PR.

Also on this:

I don't believe that anything should go to CPU though. That should only happen if the target is on CPU.

we hit CPU only for interpolation (which isn't yet supported with Reactant), which is what we fall back to if we decide that broadcasting / copying will not work.

Thanks Greg!

I reverted to set! and it's not yet working here end-to-end, I'll investigate a bit later.

Ah sorry, stupid me, I only added the branch in the main and normal test env of Terrarium and not in the Reactant test env. It does actually work.

bgroenks96

Looks overall good, just look at the set! thing. That would be problematic if Reactant can't handle that.

maximilian-gelbrecht · 2026-08-04T08:32:56Z

Oh I thought this would already pass but it turns out there is another upstream Oceananigans Reactant issue with boundary conditions that we hit here. I converted this to draft for now.

The main problem is the 2D indexing of Fields as it seems.

I need to sit down and look at this properly to understand exactly what we need and maybe fix the issue upstream in Oceananigans.

Needs further investigation

bgroenks96 · 2026-08-04T09:41:17Z

+  StableHLO raise pass cannot lift. An unrolled loop raises cleanly.
+- **Reverse-mode AD**: a fixed trip count avoids a dynamically-sized tape.
+"""
+struct NewtonSolver{NF, iterations} end


Why? We already had a Newton solver before and got rid of it in favor of RootSolvers.jl.

Can't you just disable the convergence test for the RootSolvers one?

Also, this is again something that should be fixed with a custom rule... the AD system should not be differentiating through the iterations.

No custom rules with Reactant, but you are probably right with RootSolvers.

Is there any documentation or roadmap anywhere about custom rules? I can't find anything, and I find this really concerning.

I am not terribly concerned around it as even the worst case of just differentiating through the maximum itertation number might not be that bad. But we are also not alone with this CliMA/Oceananigans.jl#5844

@olivierbonte I also attempted it (not sure if I pushed it to my fork) and ran into the same issues. I gave up after a couple of hours. I don't think it's entirely an issue of RootSolvers.jl being incompatible; in theory Reactant can handle data-dependent loops. But I hit a point where it looked like bugs in Reactant/StableHLO and just gave up.

Alright, fixed point iteration for Reactant for now sounds reasonable!
Good to know you also gave it a shot @bgroenks96! I'll try to look at it a bit more if I find the time, but not sure If I will get the AD working. Do you think there is value in trying to get a PR merged in RootSolvers.jl that works for forward evaluations (but not yet for Enzyme)? Because that I think should work

If the changes to the code are not too substantial, I would try it. But the CliMA project is pretty explicitly disinterested in AD so I doubt they will accept a major rewrite for Reactant.

If you can reduce it to a Reactant MWE if something isn't working that would be great as well though, we should absolutely post it in their repo. Usually responses are quick.

CliMA/RootSolvers.jl#99 now works for primal evaluations. Enzyme autodiff still WIP

bgroenks96 · 2026-08-04T12:32:25Z

The main problem is the 2D indexing of Fields as it seems.

It would not be so objectionable to do field[i, j, 1] rather than field[i, j] everywhere. We already have to do this anyway when using setindex!.

maximilian-gelbrecht · 2026-08-04T13:02:55Z

The main problem is the 2D indexing of Fields as it seems.

It would not be so objectionable to do field[i, j, 1] rather than field[i, j] everywhere. We already have to do this anyway when using setindex!.

The culprit is in Oceananigans though CliMA/Oceananigans.jl#5846

maximilian-gelbrecht · 2026-08-10T17:36:36Z

The main problem is the 2D indexing of Fields as it seems.

It would not be so objectionable to do field[i, j, 1] rather than field[i, j] everywhere. We already have to do this anyway when using setindex!.

The culprit is in Oceananigans though CliMA/Oceananigans.jl#5846

This issue is resolved (but not yet released).

I'll pick up this PR when I find some time, maybe on the way home from JuliaCon.

…ium.jl into mg/reactant-landmodel

maximilian-gelbrecht · 2026-08-13T14:15:27Z

This works now. I would keep the fixed iteration NewtonSolver as a temporary solution for Reactant until we have something different figured out (like RootSolvers.jl compatability).

olivierbonte · 2026-08-13T14:25:48Z

@maximilian-gelbrecht have you tried the RootSolvers.jl version at CliMA/RootSolvers.jl#99? This was working okay for normal (non AD) reactant code evaluations but haven't tried it yet within Terrarium.

maximilian-gelbrecht · 2026-08-17T06:55:00Z

Further Oceananigans issues that blocked the vegetation model are addressed in CliMA/Oceananigans.jl#5880

land model reactant wo veg

133e7bb

maximilian-gelbrecht requested a review from bgroenks96 August 4, 2026 07:35

maximilian-gelbrecht commented Aug 4, 2026

View reviewed changes

Comment thread src/processes/soil/hydrology/soil_hydraulic_closures.jl

bgroenks96 reviewed Aug 4, 2026

View reviewed changes

Comment thread src/processes/soil/hydrology/soil_hydraulic_properties.jl

remove comment

7c33e5a

bgroenks96 reviewed Aug 4, 2026

View reviewed changes

bgroenks96 previously approved these changes Aug 4, 2026

View reviewed changes

pad indices

758d8d6

maximilian-gelbrecht marked this pull request as draft August 4, 2026 08:22

maximilian-gelbrecht added 3 commits August 4, 2026 10:40

nonbranching compute_hydraulics!

0fcb7da

fixed iteration newton solver

ffed50f

commit claude plan

536833e

bgroenks96 reviewed Aug 4, 2026

View reviewed changes

maximilian-gelbrecht mentioned this pull request Aug 4, 2026

getbc fails with Reactant due to 2D indexing CliMA/Oceananigans.jl#5846

Closed

maximilian-gelbrecht added 2 commits August 4, 2026 16:15

update plans

b36e278

update plan

18842ac

maximilian-gelbrecht mentioned this pull request Aug 6, 2026

Restore full GPU compat of LandModel #160

Merged

Merge branch 'main' into mg/reactant-landmodel

7620b15

maximilian-gelbrecht mentioned this pull request Aug 10, 2026

Reactant compatibility: further models and development notes #143

Open

olivierbonte mentioned this pull request Aug 11, 2026

Attempt at making RootSolvers.jl compatible with Reactant CliMA/RootSolvers.jl#99

Open

maximilian-gelbrecht added 5 commits August 13, 2026 13:59

Merge remote-tracking branch 'origin/main' into mg/reactant-landmodel

edc4ac9

Merge branch 'mg/reactant-landmodel' of github.com:TUM-PIK-ESM/Terrar…

1f53cd1

…ium.jl into mg/reactant-landmodel

@assert_kernel added

3d6ba6d

add reactantcore to the deps

7ffab6f

update macro

bf37231

remove pad_indices again

283e58a

maximilian-gelbrecht commented Aug 13, 2026

View reviewed changes

Comment thread src/utils/kernel_utils.jl

maximilian-gelbrecht added 9 commits August 13, 2026 14:41

don't use within_compile() anymore

f7f5816

remove reactantcore

6a362fd

formatting

a4abf98

update Oceananigans compat bounds

6483f35

use field_indices

8c9179d

comments

0365122

use @assert_kernel

779d255

update plan

ce4fda3

update plan again, more comments

faa6470

maximilian-gelbrecht marked this pull request as ready for review August 13, 2026 14:14

glwagner mentioned this pull request Aug 13, 2026

Name the copy-versus-interpolate cases and share them with Reactant CliMA/Oceananigans.jl#5878

Open

maximilian-gelbrecht added 7 commits August 14, 2026 10:29

formatting

c7f2bac

revert to set!

dcae394

use PR branch of Oceananigans

c3cb8ad

Merge branch 'main' into mg/reactant-landmodel

8e35e49

also use the branch in the test env

bf219ce

update plan and env

99c748b

update plan

1bae7c0

Uh oh!

Conversation

maximilian-gelbrecht commented Aug 4, 2026

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

maximilian-gelbrecht Aug 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bgroenks96 left a comment

Choose a reason for hiding this comment

Uh oh!

maximilian-gelbrecht commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bgroenks96 Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bgroenks96 commented Aug 4, 2026

Uh oh!

maximilian-gelbrecht commented Aug 4, 2026

Uh oh!

maximilian-gelbrecht commented Aug 10, 2026

Uh oh!

Uh oh!

maximilian-gelbrecht commented Aug 13, 2026

Uh oh!

olivierbonte commented Aug 13, 2026

Uh oh!

maximilian-gelbrecht commented Aug 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

maximilian-gelbrecht Aug 14, 2026 •

edited

Loading

maximilian-gelbrecht commented Aug 4, 2026 •

edited

Loading

bgroenks96 Aug 4, 2026 •

edited

Loading