From 94f202d9a98de6387e2abd376945c05f1ea1752f Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 3 Jan 2025 13:48:28 +0100
Subject: [PATCH 01/62] diff tests in extra folder

---
 test/differentiability/barotropic.jl        |  53 ++++
 test/differentiability/runtests.jl          |  14 +
 test/differentiability/speedy_transforms.jl | 314 +++++++++++++++++++
 test/runtests.jl                            |   5 +
 test/spectral_transform_ad_rules.jl         | 322 +-------------------
 5 files changed, 388 insertions(+), 320 deletions(-)
 create mode 100644 test/differentiability/barotropic.jl
 create mode 100644 test/differentiability/runtests.jl
 create mode 100644 test/differentiability/speedy_transforms.jl

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
new file mode 100644
index 000000000..c76ec509c
--- /dev/null
+++ b/test/differentiability/barotropic.jl
@@ -0,0 +1,53 @@
+
+@testset "Differentiability: Barotropic Model Timestepping" begin 
+    spectral_grid = SpectralGrid(trunc=31, nlayers=1)          # define resolution
+    model = BarotropicModel(; spectral_grid)   # construct model
+    simulation = initialize!(model)  
+
+    (; prognostic_variables, diagnostic_variables, model) = simulation
+    (; clock) = prognostic_variables
+
+    period = Day(10)
+    progn = prognostic_variables
+    diagn = diagnostic_variables
+
+    # CLOCK
+    SpeedyWeather.set_period!(clock, period)              # set how long to integrate for
+    SpeedyWeather.initialize!(clock, model.time_stepping) # store the start date, reset counter
+
+    (; Δt, Δt_millisec) = model.time_stepping
+
+    # SCALING: we use vorticity*radius, divergence*radius in the dynamical core
+    SpeedyWeather.scale!(progn, diagn, model.spectral_grid.radius)
+
+    # OUTPUT INITIALISATION AND STORING INITIAL CONDITIONS + FEEDBACK
+    # propagate spectral state to grid variables for initial condition output
+    (; output, feedback) = model
+    lf = 1                                  # use first leapfrog index
+    SpeedyWeather.transform!(diagn, progn, lf, model, initialize=true)
+    SpeedyWeather.initialize!(progn.particles, progn, diagn, model.particle_advection)
+    SpeedyWeather.initialize!(output, feedback, progn, diagn, model)
+    SpeedyWeather.initialize!(model.callbacks, progn, diagn, model)
+
+    # FIRST TIMESTEPS: EULER FORWARD THEN 1x LEAPFROG
+    # considered part of the model initialisation
+    SpeedyWeather.first_timesteps!(progn, diagn, model)
+
+    # only now initialise feedback for benchmark accuracy
+    SpeedyWeather.initialize!(feedback, clock, model)
+
+    d_progn = PrognosticVariables(spectral_grid)
+    d_diag = DiagnosticVariables(spectral_grid)
+
+    d_model = deepcopy(model)
+
+
+    SpeedyWeather.timestep!(progn, diagn, 2Δt, model) # calculate tendencies and leapfrog forward
+
+    autodiff(Reverse, SpeedyWeather.timestep!, Const, Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(2Δt), Duplicated(model, d_model))
+
+    # differnetiate wrt initial conditions 
+
+    # differnetiate wrt parameter 
+
+end
\ No newline at end of file
diff --git a/test/differentiability/runtests.jl b/test/differentiability/runtests.jl
new file mode 100644
index 000000000..3da7db777
--- /dev/null
+++ b/test/differentiability/runtests.jl
@@ -0,0 +1,14 @@
+# Here we test the differentiability of the model by comparing Enzyme's results to finite difference differenation 
+# These tests are relatively expensive, and also not strictly necessary to perform at every commit, so they 
+# are only part of the extended test set 
+using SpeedyWeather, EnzymeTestUtils, Enzyme, FiniteDifferences
+import EnzymeTestUtils: test_approx
+import FiniteDifferences: j′vp, grad, central_fdm
+import AbstractFFTs
+
+grid_types = [FullGaussianGrid, OctahedralGaussianGrid] # one full and one reduced grid, both Gaussian to have exact transforms 
+grid_dealiasing = [2, 3]
+fd_tests = [true, true] 
+
+include("speedy_transforms.jl")
+include("barotropic.jl")
\ No newline at end of file
diff --git a/test/differentiability/speedy_transforms.jl b/test/differentiability/speedy_transforms.jl
new file mode 100644
index 000000000..19869e12e
--- /dev/null
+++ b/test/differentiability/speedy_transforms.jl
@@ -0,0 +1,314 @@
+# Tests for SpeedyTransforms
+
+@testset "Differentiability: Complete Transform Enzyme" begin
+    # make a high level finite difference test of the whole transform
+    # can't use Enzyme or ChainRule Test tools for tests for that
+    for (i_grid, grid_type) in enumerate(grid_types)
+
+        spectral_grid = SpectralGrid(Grid=grid_type, trunc=10, nlayers=1, dealiasing=grid_dealiasing[i_grid])
+        S = SpectralTransform(spectral_grid, one_more_degree=true)
+        dS = deepcopy(S)
+
+        if fd_tests[i_grid]
+            
+            # forwards 
+            grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
+            dgrid = zero(grid)
+            specs = zeros(LowerTriangularArray{Complex{spectral_grid.NF}}, spectral_grid.trunc+2, spectral_grid.trunc+1, spectral_grid.nlayers)
+            
+            # seed
+            dspecs = zero(specs)
+            fill!(dspecs, 1+1im)
+
+            autodiff(Reverse, transform!, Const, Duplicated(specs, dspecs), Duplicated(grid, dgrid), Duplicated(S, dS))
+
+            # new seed
+            dspecs2 = zero(specs)
+            fill!(dspecs2, 1+1im)
+
+            # finite difference comparision, seeded with a one adjoint to get the direct gradient
+            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dspecs2, grid)
+            @test isapprox(dgrid, fd_jvp[1])
+
+            ## now backwards, as the input for spec we use the output of the forward transform
+
+            fill!(dspecs,0)
+            grid = zeros(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
+            dgrid = similar(grid)
+            fill!(dgrid, 1)
+
+            autodiff(Reverse, transform!, Const, Duplicated(grid, dgrid), Duplicated(specs, dspecs), Duplicated(S, dS))
+
+            # new seed 
+            dgrid2 = similar(grid)
+            fill!(dgrid2, 1)
+
+            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dgrid2, specs)
+
+            @test isapprox(dspecs, fd_jvp[1])
+        end 
+
+        # test that d S^{-1}(S(x)) / dx = dx/dx = 1 (starting in both domains) 
+        # this only holds for exact transforms, like Gaussian grids 
+
+        # start with grid (but with a truncated one)
+        function transform_identity!(x_out::AbstractGridArray{T}, x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
+            x_SH = zeros(LowerTriangularArray{Complex{T}}, S.lmax+1, S.mmax+1, S.nlayers)
+            transform!(x_SH, x, S)
+            transform!(x_out, x_SH, S)
+            return nothing
+        end 
+
+        function transform_identity(x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
+            x_copy = deepcopy(x)
+            transform_identity!(x_copy, x, S) 
+            return x_copy
+        end
+        
+        grid = rand(S.Grid{spectral_grid.NF}, S.nlat_half, S.nlayers)
+        spec = transform(grid, S)
+
+        grid = transform(spec, S)
+        grid_out = zero(grid)
+        
+        transform_identity!(grid_out, grid, S)
+        @test isapprox(grid, grid_out)
+
+        dgrid = similar(grid)
+        fill!(dgrid, 1)
+
+        dgrid_out = zero(grid_out)
+
+        autodiff(Reverse, transform_identity!, Const, Duplicated(grid_out, dgrid_out), Duplicated(grid, dgrid), Duplicated(S, dS))
+
+        @test all(isapprox.(dgrid, 1)) 
+        # TODO: previously this test was broken, with a version that directly mutates in-place. 
+        # FD yields the same non-one values though. 
+        # Not sure why. Do we use such things in our model? 
+        #
+        #function transform_identity!(x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
+        #   x_SH = zeros(LowerTriangularArray{Complex{T}}, S.lmax+1, S.mmax+1, S.nlayers)
+        #   transform!(x_SH, x, S)
+        #   transform!(x, x_SH, S)
+        #   return nothing
+        #end 
+        # The FD comparision passes, but it takes a long time to compute, so it's commented out. 
+        #dgrid2 = similar(grid)
+        #fill!(dgrid2, 1)
+        #fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_identity(x, S), dgrid2, grid)
+        #@test isapprox(dgrid, fd_jvp[1], rtol=0.01)
+
+        # now start with spectral space, exclude for other grid because of https://github.com/SpeedyWeather/SpeedyWeather.jl/issues/626
+        if fd_tests[i_grid]
+
+            function transform_identity!(x::LowerTriangularArray{Complex{T}}, S::SpectralTransform{T}) where T
+                x_grid = zeros(S.Grid{T}, S.nlat_half, S.nlayers)
+                transform!(x_grid, x, S)
+                transform!(x, x_grid, S)
+                return nothing
+            end 
+
+            spec = transform(grid, S)
+            spec_copy = deepcopy(spec)
+            transform_identity!(spec, S)
+            @test isapprox(spec, spec_copy)
+
+            dspec = similar(spec)
+            fill!(dspec, 1+im)
+
+            autodiff(Reverse, transform_identity!, Const, Duplicated(spec, dspec), Duplicated(S, dS))
+
+            @test all(all.([isapprox.(dspec[il,1,:], 1) for il in 1:S.lmax+1])) # m = 0 => Im = 0
+
+            for i in eachmatrix(dspec)
+                @test all(isapprox.(dspec[:,i][S.lmax+2:end], 1+im)) 
+            end 
+        end
+    end 
+end
+
+# Tests for all other spectral gradient functions
+# We test that gradients are non-zero and identical with their finite difference
+# When easiliy possible we check with the analytical formula as well 
+@testset "Differentiability: Spectral Gradients Enzyme" begin 
+    for (i_grid, grid_type) in enumerate(grid_types)
+
+        if fd_tests[i_grid]
+
+            spectral_grid = SpectralGrid(Grid=grid_type, trunc=10, nlayers=1, dealiasing=grid_dealiasing[i_grid])
+            S = SpectralTransform(spectral_grid, one_more_degree=true)
+            dS = deepcopy(S)
+
+            u_grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
+            v_grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
+
+            u = transform(u_grid, S)
+            v = transform(v_grid, S)
+            du = zero(u)
+            dv = zero(v)
+
+            cu = zero(u)
+            dcu = zero(u)
+            fill!(dcu, 1)
+
+            # curl test
+            autodiff(Reverse, curl!, Const, Duplicated(cu, dcu), Duplicated(u, du), Duplicated(v, dv), Duplicated(S, dS))
+            
+            # we know  the gradient of the divergence wrt v is easy: im*m
+            # so with a seed of 1, we should get for dv: im*m * 1 = im*m 
+            # See https://speedyweather.github.io/SpeedyWeather.jl/dev/spectral_transform/
+            # let's check it
+            # TO-DO: why the other sign? but it's the same for Finite Differences
+            # It's because it's the adjoint (')? And this matters here for complex numbers (see e.g. FiniteDifferences.jl examples for j'vp)
+            # To-Do: double check that
+            for i=1:dv.n
+                @test all(Array(dv[:,1])[i:dv.m-1,i] .≈ complex(0,-(i-1)))
+            end  
+            @test sum(du) != 0 # nonzero gradient
+
+            # new seed
+            dcu2 = zero(dcu)
+            fill!(dcu2, 1)
+
+            # finite difference comparision, seeded with a one adjoint to get the direct gradient
+            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> curl(x[1],x[2], S), dcu2, (u, v))
+            @test isapprox(du, fd_jvp[1][1], rtol=1e-6)
+            @test isapprox(dv, fd_jvp[1][2], rtol=1e-6)
+
+            # div test
+            u = transform(u_grid, S)
+            v = transform(v_grid, S) 
+            du = zero(u)
+            dv = zero(v)
+            div = zero(u)
+            ddiv = zero(u)
+            fill!(ddiv, 1)
+
+            autodiff(Reverse, divergence!, Const, Duplicated(div, ddiv), Duplicated(u, du), Duplicated(v, dv), Duplicated(S, dS))
+
+            # we know the gradient of the divergence wrt u is easy: im*m 
+            # See https://speedyweather.github.io/SpeedyWeather.jl/dev/spectral_transform/
+            # let's check it 
+            # To-Do: why the minus sign? 
+            # It's because it's the adjoint (')? And this matters here for complex numbers
+            # To-Do: double check that
+            for i=1:du.n
+                @test all(Array(du[:,1])[i:du.m-1,i] .≈ complex(i-1,-(i-1)))
+            end 
+
+            ddiv2 = zero(ddiv)
+            fill!(ddiv2, 1)
+
+            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> divergence(x[1],x[2], S), ddiv2, (u, v))
+            @test isapprox(du, fd_jvp[1][1])
+            @test isapprox(dv, fd_jvp[1][2])
+            @test sum(du) != 0 # nonzero gradient
+            @test sum(dv) != 0 # nonzero gradient
+
+            # UV_from_vor! 
+
+            u = zero(u)
+            du = fill!(du, 1+1im)
+
+            v = zero(v)
+            dv = fill!(dv, 1+1im)
+
+            vor_grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
+            vor = transform(vor_grid, S)
+            dvor = zero(vor)
+
+            autodiff(Reverse, SpeedyWeather.SpeedyTransforms.UV_from_vor!, Const, Duplicated(u, du), Duplicated(v, dv), Duplicated(vor, dvor), Duplicated(S, dS))
+
+            function uvfvor(vor, S)
+                u = zero(vor)
+                v = zero(vor)
+                SpeedyWeather.SpeedyTransforms.UV_from_vor!(u, v, vor, S)
+                return cat(u, v, dims=2)
+            end 
+            
+            uv_input = cat(u, v, dims=2)
+            duv_input = zero(uv_input)
+            duv_input = fill!(duv_input, 1+im)
+
+            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> uvfvor(x, S), duv_input, vor)
+            @test isapprox(dvor, fd_jvp[1])
+            @test sum(dvor) != 0 # nonzero gradient
+
+            # UV_from_vordiv! 
+            u = zero(u)
+            du = fill!(du, 1+1im)
+
+            v = zero(v)
+            dv = fill!(dv, 1+1im)
+
+            vor_grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
+            vor = transform(vor_grid, S)
+            dvor = zero(vor)
+
+            div_grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
+            div = transform(vor_grid, S)
+            ddiv = zero(vor)
+
+            autodiff(Reverse, SpeedyWeather.SpeedyTransforms.UV_from_vordiv!, Const, Duplicated(u, du), Duplicated(v, dv), Duplicated(vor, dvor), Duplicated(div, ddiv), Duplicated(S, dS))
+
+            function uvfromvordiv(vor, div, S)
+                u = zero(vor)
+                v = zero(vor)
+                SpeedyWeather.SpeedyTransforms.UV_from_vordiv!(u, v, vor, div, S)
+                return cat(u, v, dims=2)
+            end 
+            
+            uv_input = zero(uv_input)
+            duv_input = fill!(duv_input, 1+im)
+
+            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x-> uvfromvordiv(x[1], x[2], S), duv_input, (vor, div))
+            @test isapprox(dvor, fd_jvp[1][1][:,1]) 
+            @test isapprox(ddiv, fd_jvp[1][2][:,1])
+            @test sum(dvor) != 0 # nonzero gradient
+            @test sum(ddiv) != 0 # nonzero gradient
+
+            # ∇²
+            dvor = zero(vor)
+            res_∇ = zero(vor)
+            dres_∇ = zero(res_∇)
+            fill!(dres_∇, 1+im)
+
+            autodiff(Reverse, SpeedyWeather.SpeedyTransforms.∇²!, Const, Duplicated(res_∇, dres_∇), Duplicated(vor, dvor), Duplicated(S, dS))
+
+            dres_∇2 = zero(res_∇)
+            fill!(dres_∇2, 1+im)
+
+            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x-> ∇²(x, S), dres_∇2, vor)
+            @test sum(dvor) != 0 # non-zero 
+            @test isapprox(dvor, fd_jvp[1]) # and identical with FD
+
+            # test with the eigenvalues saved in S, result should just be seed * eigenvalues
+            for i=1:(vor.m-1)
+                @test all(isapprox.(Array(dvor[:,1])[i,1:i], S.eigenvalues[i] * (1+im)))
+            end 
+
+            # ∇
+            zonal_gradient = zero(vor)
+            dzonal_gradient = zero(vor)
+            fill!(dzonal_gradient, 1+im)
+
+            merid_gradient = zero(vor)
+            dmerid_gradient = zero(vor)
+            fill!(dmerid_gradient, 1+im)
+
+            dvor = zero(vor)
+            autodiff(Reverse, SpeedyWeather.SpeedyTransforms.∇!, Const, Duplicated(zonal_gradient, dzonal_gradient), Duplicated(merid_gradient, dmerid_gradient), Duplicated(vor, dvor), Duplicated(S, dS))
+
+            dmerid_gradient2 = zero(dmerid_gradient)
+            fill!(dmerid_gradient2, 1+im)
+
+            dzonal_gradient2 = zero(dzonal_gradient)
+            fill!(dzonal_gradient2, 1+im)
+
+            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x-> ∇(x, S), (dmerid_gradient2, dzonal_gradient2), vor)
+            @test sum(dvor) != # nonzero 
+            @test isapprox(dvor, fd_jvp[1]) # and identical with FD
+        end 
+    end 
+end 
+ 
diff --git a/test/runtests.jl b/test/runtests.jl
index 76f6e7925..ef561f38d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -58,5 +58,10 @@ include("run_speedy.jl")
 include("callbacks.jl")
 include("schedule.jl")
 
+# DIFFERENTIABILITY 
+if FLAG_EXTENDED_TESTS
+    include("differentiability/runtests.jl")
+end
+
 # OUTPUT 
 include("netcdf_output.jl")
\ No newline at end of file
diff --git a/test/spectral_transform_ad_rules.jl b/test/spectral_transform_ad_rules.jl
index 2562764e7..3a68c0062 100644
--- a/test/spectral_transform_ad_rules.jl
+++ b/test/spectral_transform_ad_rules.jl
@@ -1,7 +1,5 @@
-using SpeedyWeather
-using EnzymeTestUtils, Enzyme
-import EnzymeTestUtils: test_approx
-using FiniteDifferences
+using EnzymeTestUtils, Enzyme, FiniteDifferences
+import EnzymeTestUtils: test_approx 
 import FiniteDifferences: j′vp, grad, central_fdm
 import AbstractFFTs
 
@@ -55,322 +53,6 @@ end
             end
         end
     end 
-
-    if FLAG_EXTENDED_TESTS # part of the extented tests, not tested in regular CI 
-
-        @testset "Complete Transform Enzyme" begin
-            # make a high level finite difference test of the whole transform
-            # can't use Enzyme or ChainRule Test tools for tests for that
-            for (i_grid, grid_type) in enumerate(grid_types)
-
-                spectral_grid = SpectralGrid(Grid=grid_type, trunc=10, nlayers=1, dealiasing=grid_dealiasing[i_grid])
-                S = SpectralTransform(spectral_grid, one_more_degree=true)
-                dS = deepcopy(S)
-
-                if fd_tests[i_grid]
-                    
-                    # forwards 
-                    grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
-                    dgrid = zero(grid)
-                    specs = zeros(LowerTriangularArray{Complex{spectral_grid.NF}}, spectral_grid.trunc+2, spectral_grid.trunc+1, spectral_grid.nlayers)
-                    
-                    # seed
-                    dspecs = zero(specs)
-                    fill!(dspecs, 1+1im)
-
-                    autodiff(Reverse, transform!, Const, Duplicated(specs, dspecs), Duplicated(grid, dgrid), Duplicated(S, dS))
-
-                    # new seed
-                    dspecs2 = zero(specs)
-                    fill!(dspecs2, 1+1im)
-
-                    # finite difference comparision, seeded with a one adjoint to get the direct gradient
-                    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dspecs2, grid)
-                    @test isapprox(dgrid, fd_jvp[1])
-
-                    ## now backwards, as the input for spec we use the output of the forward transform
-
-                    fill!(dspecs,0)
-                    grid = zeros(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
-                    dgrid = similar(grid)
-                    fill!(dgrid, 1)
-
-                    autodiff(Reverse, transform!, Const, Duplicated(grid, dgrid), Duplicated(specs, dspecs), Duplicated(S, dS))
-
-                    # new seed 
-                    dgrid2 = similar(grid)
-                    fill!(dgrid2, 1)
-
-                    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dgrid2, specs)
-
-                    @test isapprox(dspecs, fd_jvp[1])
-                end 
-
-                # test that d S^{-1}(S(x)) / dx = dx/dx = 1 (starting in both domains) 
-                # this only holds for exact transforms, like Gaussian grids 
-
-                # start with grid (but with a truncated one)
-                function transform_identity!(x_out::AbstractGridArray{T}, x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
-                    x_SH = zeros(LowerTriangularArray{Complex{T}}, S.lmax+1, S.mmax+1, S.nlayers)
-                    transform!(x_SH, x, S)
-                    transform!(x_out, x_SH, S)
-                    return nothing
-                end 
-
-                function transform_identity(x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
-                    x_copy = deepcopy(x)
-                    transform_identity!(x_copy, x, S) 
-                    return x_copy
-                end
-                
-                grid = rand(S.Grid{spectral_grid.NF}, S.nlat_half, S.nlayers)
-                spec = transform(grid, S)
-
-                grid = transform(spec, S)
-                grid_out = zero(grid)
-                
-                transform_identity!(grid_out, grid, S)
-                @test isapprox(grid, grid_out)
-
-                dgrid = similar(grid)
-                fill!(dgrid, 1)
-
-                dgrid_out = zero(grid_out)
-
-                autodiff(Reverse, transform_identity!, Const, Duplicated(grid_out, dgrid_out), Duplicated(grid, dgrid), Duplicated(S, dS))
-
-                @test all(isapprox.(dgrid, 1)) 
-                # TODO: previously this test was broken, with a version that directly mutates in-place. 
-                # FD yields the same non-one values though. 
-                # Not sure why. Do we use such things in our model? 
-                #
-                #function transform_identity!(x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
-                #   x_SH = zeros(LowerTriangularArray{Complex{T}}, S.lmax+1, S.mmax+1, S.nlayers)
-                #   transform!(x_SH, x, S)
-                #   transform!(x, x_SH, S)
-                #   return nothing
-                #end 
-                # The FD comparision passes, but it takes a long time to compute, so it's commented out. 
-                #dgrid2 = similar(grid)
-                #fill!(dgrid2, 1)
-                #fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_identity(x, S), dgrid2, grid)
-                #@test isapprox(dgrid, fd_jvp[1], rtol=0.01)
-
-                # now start with spectral space, exclude for other grid because of https://github.com/SpeedyWeather/SpeedyWeather.jl/issues/626
-                if fd_tests[i_grid]
-
-                    function transform_identity!(x::LowerTriangularArray{Complex{T}}, S::SpectralTransform{T}) where T
-                        x_grid = zeros(S.Grid{T}, S.nlat_half, S.nlayers)
-                        transform!(x_grid, x, S)
-                        transform!(x, x_grid, S)
-                        return nothing
-                    end 
-
-                    spec = transform(grid, S)
-                    spec_copy = deepcopy(spec)
-                    transform_identity!(spec, S)
-                    @test isapprox(spec, spec_copy)
-
-                    dspec = similar(spec)
-                    fill!(dspec, 1+im)
-
-                    autodiff(Reverse, transform_identity!, Const, Duplicated(spec, dspec), Duplicated(S, dS))
-
-                    @test all(all.([isapprox.(dspec[il,1,:], 1) for il in 1:S.lmax+1])) # m = 0 => Im = 0
-
-                    for i in eachmatrix(dspec)
-                        @test all(isapprox.(dspec[:,i][S.lmax+2:end], 1+im)) 
-                    end 
-                end
-            end 
-        end
-
-        # Tests for all other spectral gradient functions
-        # We test that gradients are non-zero and identical with their finite difference
-        # When easiliy possible we check with the analytical formula as well 
-        @testset "Spectral Gradient Enzyme" begin 
-            for (i_grid, grid_type) in enumerate(grid_types)
-
-                if fd_tests[i_grid]
-
-                    spectral_grid = SpectralGrid(Grid=grid_type, trunc=10, nlayers=1, dealiasing=grid_dealiasing[i_grid])
-                    S = SpectralTransform(spectral_grid, one_more_degree=true)
-                    dS = deepcopy(S)
-
-                    u_grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
-                    v_grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
-
-                    u = transform(u_grid, S)
-                    v = transform(v_grid, S)
-                    du = zero(u)
-                    dv = zero(v)
-
-                    cu = zero(u)
-                    dcu = zero(u)
-                    fill!(dcu, 1)
-
-                    # curl test
-                    autodiff(Reverse, curl!, Const, Duplicated(cu, dcu), Duplicated(u, du), Duplicated(v, dv), Duplicated(S, dS))
-                    
-                    # we know  the gradient of the divergence wrt v is easy: im*m
-                    # so with a seed of 1, we should get for dv: im*m * 1 = im*m 
-                    # See https://speedyweather.github.io/SpeedyWeather.jl/dev/spectral_transform/
-                    # let's check it
-                    # TO-DO: why the other sign? but it's the same for Finite Differences
-                    # It's because it's the adjoint (')? And this matters here for complex numbers (see e.g. FiniteDifferences.jl examples for j'vp)
-                    # To-Do: double check that
-                    for i=1:dv.n
-                        @test all(Array(dv[:,1])[i:dv.m-1,i] .≈ complex(0,-(i-1)))
-                    end  
-                    @test sum(du) != 0 # nonzero gradient
-
-                    # new seed
-                    dcu2 = zero(dcu)
-                    fill!(dcu2, 1)
-
-                    # finite difference comparision, seeded with a one adjoint to get the direct gradient
-                    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> curl(x[1],x[2], S), dcu2, (u, v))
-                    @test isapprox(du, fd_jvp[1][1], rtol=1e-6)
-                    @test isapprox(dv, fd_jvp[1][2], rtol=1e-6)
-
-                    # div test
-                    u = transform(u_grid, S)
-                    v = transform(v_grid, S) 
-                    du = zero(u)
-                    dv = zero(v)
-                    div = zero(u)
-                    ddiv = zero(u)
-                    fill!(ddiv, 1)
-
-                    autodiff(Reverse, divergence!, Const, Duplicated(div, ddiv), Duplicated(u, du), Duplicated(v, dv), Duplicated(S, dS))
-
-                    # we know the gradient of the divergence wrt u is easy: im*m 
-                    # See https://speedyweather.github.io/SpeedyWeather.jl/dev/spectral_transform/
-                    # let's check it 
-                    # To-Do: why the minus sign? 
-                    # It's because it's the adjoint (')? And this matters here for complex numbers
-                    # To-Do: double check that
-                    for i=1:du.n
-                        @test all(Array(du[:,1])[i:du.m-1,i] .≈ complex(i-1,-(i-1)))
-                    end 
-
-                    ddiv2 = zero(ddiv)
-                    fill!(ddiv2, 1)
-
-                    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> divergence(x[1],x[2], S), ddiv2, (u, v))
-                    @test isapprox(du, fd_jvp[1][1])
-                    @test isapprox(dv, fd_jvp[1][2])
-                    @test sum(du) != 0 # nonzero gradient
-                    @test sum(dv) != 0 # nonzero gradient
-
-                    # UV_from_vor! 
-
-                    u = zero(u)
-                    du = fill!(du, 1+1im)
-
-                    v = zero(v)
-                    dv = fill!(dv, 1+1im)
-
-                    vor_grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
-                    vor = transform(vor_grid, S)
-                    dvor = zero(vor)
-
-                    autodiff(Reverse, SpeedyWeather.SpeedyTransforms.UV_from_vor!, Const, Duplicated(u, du), Duplicated(v, dv), Duplicated(vor, dvor), Duplicated(S, dS))
-
-                    function uvfvor(vor, S)
-                        u = zero(vor)
-                        v = zero(vor)
-                        SpeedyWeather.SpeedyTransforms.UV_from_vor!(u, v, vor, S)
-                        return cat(u, v, dims=2)
-                    end 
-                    
-                    uv_input = cat(u, v, dims=2)
-                    duv_input = zero(uv_input)
-                    duv_input = fill!(duv_input, 1+im)
-
-                    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> uvfvor(x, S), duv_input, vor)
-                    @test isapprox(dvor, fd_jvp[1])
-                    @test sum(dvor) != 0 # nonzero gradient
-
-                    # UV_from_vordiv! 
-                    u = zero(u)
-                    du = fill!(du, 1+1im)
-
-                    v = zero(v)
-                    dv = fill!(dv, 1+1im)
-
-                    vor_grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
-                    vor = transform(vor_grid, S)
-                    dvor = zero(vor)
-
-                    div_grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
-                    div = transform(vor_grid, S)
-                    ddiv = zero(vor)
-
-                    autodiff(Reverse, SpeedyWeather.SpeedyTransforms.UV_from_vordiv!, Const, Duplicated(u, du), Duplicated(v, dv), Duplicated(vor, dvor), Duplicated(div, ddiv), Duplicated(S, dS))
-
-                    function uvfromvordiv(vor, div, S)
-                        u = zero(vor)
-                        v = zero(vor)
-                        SpeedyWeather.SpeedyTransforms.UV_from_vordiv!(u, v, vor, div, S)
-                        return cat(u, v, dims=2)
-                    end 
-                    
-                    uv_input = zero(uv_input)
-                    duv_input = fill!(duv_input, 1+im)
-
-                    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x-> uvfromvordiv(x[1], x[2], S), duv_input, (vor, div))
-                    @test isapprox(dvor, fd_jvp[1][1][:,1]) 
-                    @test isapprox(ddiv, fd_jvp[1][2][:,1])
-                    @test sum(dvor) != 0 # nonzero gradient
-                    @test sum(ddiv) != 0 # nonzero gradient
-
-                    # ∇²
-                    dvor = zero(vor)
-                    res_∇ = zero(vor)
-                    dres_∇ = zero(res_∇)
-                    fill!(dres_∇, 1+im)
-
-                    autodiff(Reverse, SpeedyWeather.SpeedyTransforms.∇²!, Const, Duplicated(res_∇, dres_∇), Duplicated(vor, dvor), Duplicated(S, dS))
-
-                    dres_∇2 = zero(res_∇)
-                    fill!(dres_∇2, 1+im)
-
-                    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x-> ∇²(x, S), dres_∇2, vor)
-                    @test sum(dvor) != 0 # non-zero 
-                    @test isapprox(dvor, fd_jvp[1]) # and identical with FD
-
-                    # test with the eigenvalues saved in S, result should just be seed * eigenvalues
-                    for i=1:(vor.m-1)
-                        @test all(isapprox.(Array(dvor[:,1])[i,1:i], S.eigenvalues[i] * (1+im)))
-                    end 
-
-                    # ∇
-                    zonal_gradient = zero(vor)
-                    dzonal_gradient = zero(vor)
-                    fill!(dzonal_gradient, 1+im)
-
-                    merid_gradient = zero(vor)
-                    dmerid_gradient = zero(vor)
-                    fill!(dmerid_gradient, 1+im)
-
-                    dvor = zero(vor)
-                    autodiff(Reverse, SpeedyWeather.SpeedyTransforms.∇!, Const, Duplicated(zonal_gradient, dzonal_gradient), Duplicated(merid_gradient, dmerid_gradient), Duplicated(vor, dvor), Duplicated(S, dS))
-
-                    dmerid_gradient2 = zero(dmerid_gradient)
-                    fill!(dmerid_gradient2, 1+im)
-
-                    dzonal_gradient2 = zero(dzonal_gradient)
-                    fill!(dzonal_gradient2, 1+im)
-
-                    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x-> ∇(x, S), (dmerid_gradient2, dzonal_gradient2), vor)
-                    @test sum(dvor) != # nonzero 
-                    @test isapprox(dvor, fd_jvp[1]) # and identical with FD
-                end 
-            end 
-        end 
-    end 
-
     @testset "Complete Transform ChainRules" begin 
         # WIP
     end

From 986086f74de9c734f8e452ac35db913f1fce9532 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Tue, 7 Jan 2025 10:36:05 +0100
Subject: [PATCH 02/62] more WIP

---
 test/differentiability/barotropic.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index c76ec509c..d690e5f01 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -46,8 +46,10 @@
 
     autodiff(Reverse, SpeedyWeather.timestep!, Const, Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(2Δt), Duplicated(model, d_model))
 
-    # differnetiate wrt initial conditions 
+    # differnetiate wrt initial conditions / previous state
+    # write this as functions (progn_old, diagn_old, 2\Delta t, model -> progn, diagn)
 
     # differnetiate wrt parameter 
+    # write this as function (model, progn, diagn, 2\Delta t) -> progn_new
 
 end
\ No newline at end of file

From 90c95b1e7e4bcfd34de524c71f54a09e7bbccd88 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Thu, 9 Jan 2025 17:36:00 +0100
Subject: [PATCH 03/62] new timestepping

---
 test/differentiability/barotropic.jl | 33 +++-------------------------
 1 file changed, 3 insertions(+), 30 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index d690e5f01..5dfcfbff2 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -3,46 +3,19 @@
     spectral_grid = SpectralGrid(trunc=31, nlayers=1)          # define resolution
     model = BarotropicModel(; spectral_grid)   # construct model
     simulation = initialize!(model)  
+    initialize!(simulation)
 
     (; prognostic_variables, diagnostic_variables, model) = simulation
-    (; clock) = prognostic_variables
+    (; Δt, Δt_millisec) = model.time_stepping
 
-    period = Day(10)
     progn = prognostic_variables
     diagn = diagnostic_variables
 
-    # CLOCK
-    SpeedyWeather.set_period!(clock, period)              # set how long to integrate for
-    SpeedyWeather.initialize!(clock, model.time_stepping) # store the start date, reset counter
-
-    (; Δt, Δt_millisec) = model.time_stepping
-
-    # SCALING: we use vorticity*radius, divergence*radius in the dynamical core
-    SpeedyWeather.scale!(progn, diagn, model.spectral_grid.radius)
-
-    # OUTPUT INITIALISATION AND STORING INITIAL CONDITIONS + FEEDBACK
-    # propagate spectral state to grid variables for initial condition output
-    (; output, feedback) = model
-    lf = 1                                  # use first leapfrog index
-    SpeedyWeather.transform!(diagn, progn, lf, model, initialize=true)
-    SpeedyWeather.initialize!(progn.particles, progn, diagn, model.particle_advection)
-    SpeedyWeather.initialize!(output, feedback, progn, diagn, model)
-    SpeedyWeather.initialize!(model.callbacks, progn, diagn, model)
-
-    # FIRST TIMESTEPS: EULER FORWARD THEN 1x LEAPFROG
-    # considered part of the model initialisation
-    SpeedyWeather.first_timesteps!(progn, diagn, model)
-
-    # only now initialise feedback for benchmark accuracy
-    SpeedyWeather.initialize!(feedback, clock, model)
-
     d_progn = PrognosticVariables(spectral_grid)
     d_diag = DiagnosticVariables(spectral_grid)
-
     d_model = deepcopy(model)
 
-
-    SpeedyWeather.timestep!(progn, diagn, 2Δt, model) # calculate tendencies and leapfrog forward
+    #SpeedyWeather.timestep!(progn, diagn, 2Δt, model) # calculate tendencies and leapfrog forward
 
     autodiff(Reverse, SpeedyWeather.timestep!, Const, Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(2Δt), Duplicated(model, d_model))
 

From 757dfaa72ecf5e071910ae3d1fc7e13aea08ee8e Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Thu, 9 Jan 2025 18:16:22 +0100
Subject: [PATCH 04/62] changelog diff tests timestepping

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f821587d..56fb5dda7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Differentiability tests for timestepping added [#656](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/656)
 - Bug: get_vertices for full grids had southpole at 90˚N [#654](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/654)
 - Group initialize!, timestep!, finalize! of main time loop [#652](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/652)
 - Coordinates defined in degrees lond, latd, not colat, lon [#647](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/647)

From b3545169e810a3a6b6cc08cc69be00e562544829 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Thu, 9 Jan 2025 20:43:35 +0100
Subject: [PATCH 05/62] more wip

---
 test/differentiability/barotropic.jl     | 1 +
 test/differentiability/runtests.jl       | 2 ++
 test/differentiability/timestep_utils.jl | 8 ++++++++
 3 files changed, 11 insertions(+)
 create mode 100644 test/differentiability/timestep_utils.jl

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index 5dfcfbff2..c17e11519 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -22,6 +22,7 @@
     # differnetiate wrt initial conditions / previous state
     # write this as functions (progn_old, diagn_old, 2\Delta t, model -> progn, diagn)
 
+
     # differnetiate wrt parameter 
     # write this as function (model, progn, diagn, 2\Delta t) -> progn_new
 
diff --git a/test/differentiability/runtests.jl b/test/differentiability/runtests.jl
index 3da7db777..c0cddb42c 100644
--- a/test/differentiability/runtests.jl
+++ b/test/differentiability/runtests.jl
@@ -11,4 +11,6 @@ grid_dealiasing = [2, 3]
 fd_tests = [true, true] 
 
 include("speedy_transforms.jl")
+
+include("timestep_utils.jl")
 include("barotropic.jl")
\ No newline at end of file
diff --git a/test/differentiability/timestep_utils.jl b/test/differentiability/timestep_utils.jl
new file mode 100644
index 000000000..98bae5b0b
--- /dev/null
+++ b/test/differentiability/timestep_utils.jl
@@ -0,0 +1,8 @@
+function timestep_ic!(progn_new, diagn_new, progn_old, diagn_old, dt, model)
+
+
+end 
+
+function timestep_p!() 
+
+end 
\ No newline at end of file

From 8935181313b7e0b29d128e37387162df2f4064e8 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 10 Jan 2025 15:12:45 +0100
Subject: [PATCH 06/62] prognosticvariables one, fill, zero

---
 src/dynamics/prognostic_variables.jl | 66 +++++++++++++++++++++++++++-
 test/prognostic_variables.jl         | 57 ++++++++++++++++++++++++
 test/runtests.jl                     |  1 +
 3 files changed, 123 insertions(+), 1 deletion(-)
 create mode 100644 test/prognostic_variables.jl

diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl
index 76598f407..6556bbc42 100644
--- a/src/dynamics/prognostic_variables.jl
+++ b/src/dynamics/prognostic_variables.jl
@@ -193,7 +193,10 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
     progn_new.land.soil_moisture_layer1 .= progn_old.land.soil_moisture_layer1
     progn_new.land.soil_moisture_layer2 .= progn_old.land.soil_moisture_layer2
 
-    # TODO copy over tracers
+    # copy over tracers
+    for (key, value) in progn_old.tracers
+        progn_old.tracers[key] = value
+    end 
 
     # copy largest subset of particles
     if length(progn_new.particles) != length(progn_old.particles)
@@ -206,12 +209,73 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
         progn_new.particles .= progn_old.particles
     end
 
+    progn_new.random_pattern .= progn_old.random_pattern
+
     progn_new.clock.time = progn_old.clock.time
     progn_new.scale[] = progn_old.scale[]
 
     return progn_new
 end
 
+function Base.zero(progn::PrognosticVariables{NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}) where {NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}
+    (; trunc, nlat_half, nlayers, nparticles) = progn
+    
+    # initialize regular progn variables 
+    progn_new = PrognosticVariables{NF, ArrayType, nsteps,
+        SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}(;
+            trunc, nlat_half, nlayers, nparticles,
+        )
+
+    # add tracers with zero 
+    for (key, value) in progn.tracers 
+        progn_new.tracers[key] = ntuple(i -> zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers), nsteps)
+    end 
+
+    # use the same scale 
+    progn_new.scale[] = progn.scale[]
+
+    return progn_new
+end 
+
+function Base.fill!(progn::PrognosticVariables{NF}, value::Number) where NF
+
+    value_NF = NF(value)
+
+    for i in eachindex(progn.vor)   # each leapfrog time step
+        progn.vor[i] .= value_NF
+        progn.div[i] .= value_NF
+        progn.temp[i] .= value_NF
+        progn.humid[i] .= value_NF
+        progn.pres[i] .= value_NF
+    end
+
+    # ocean
+    progn.ocean.sea_surface_temperature .= value_NF
+    progn.ocean.sea_ice_concentration .= value_NF
+    # land
+    progn.land.land_surface_temperature .= value_NF
+    progn.land.snow_depth .= value_NF
+    progn.land.soil_moisture_layer1 .= value_NF
+    progn.land.soil_moisture_layer2 .= value_NF
+
+    # fill tracers
+    for (key, value) in progn.tracers 
+        for value_i in value # istep of nsteps tuple 
+            value_i .= value_NF
+        end 
+    end 
+
+    # particles are ignored for the fill
+
+    return progn
+end 
+
+function Base.one(progn::PrognosticVariables)
+    zero_progn = zero(progn)
+    fill!(zero_progn, 1)
+    return zero_progn
+end 
+
 """$(TYPEDSIGNATURES)
 Add `tracers` to the prognostic variables `progn` in `progn.tracers::Dict`."""
 function add!(progn::PrognosticVariables{NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D}, tracers::Tracer...
diff --git a/test/prognostic_variables.jl b/test/prognostic_variables.jl
new file mode 100644
index 000000000..f5a830ec6
--- /dev/null
+++ b/test/prognostic_variables.jl
@@ -0,0 +1,57 @@
+@testset "PrognosticVariables initialize, zero, fill!, one" begin
+    
+    NF = Float32
+    nlayers = 2 
+    spectral_grid = SpectralGrid(; NF, nlayers, Grid=FullGaussianGrid)
+    model = PrimitiveWetModel(spectral_grid)
+    add!(model, Tracer(:abc))
+    simulation = initialize!(model)
+    
+    # evolve a bit to have nonzero elements 
+    run!(simulation, period=Day(1), output=false)
+
+    # zero test 
+    progn = simulation.prognostic_variables
+
+    progn_new = zero(progn)
+
+    for i in eachindex(progn_new.vor)
+        @test all(progn_new.vor[i] .== zero(NF))
+        @test all(progn_new.div[i] .== zero(NF))
+        @test all(progn_new.temp[i] .== zero(NF))
+        @test all(progn_new.humid[i] .== zero(NF))
+        @test all(progn_new.pres[i] .== zero(NF))
+    end
+
+    @test keys(progn_new.tracers) == keys(progn.tracers)
+
+    # copy test 
+    copy!(progn_new, progn)
+
+    # NaN != NaN, and there's NaNs in the ocean and land, that's why we can't directly test for equality 
+    # we don't test really all fields, it would just repeat the code that's already there in the main part
+    for i in eachindex(progn_new.vor)
+        @test all(progn_new.vor[i] .== progn.vor[i])
+        @test all(progn_new.div[i] .== progn.div[i])
+        @test all(progn_new.temp[i] .== progn.temp[i])
+        @test all(progn_new.humid[i] .== progn.humid[i])
+        @test all(progn_new.pres[i] .== progn.pres[i])
+    end
+
+    # fill! test 
+    fill!(progn_new, 1)
+
+    for i in eachindex(progn_new.vor)
+        @test all(progn_new.vor[i] .== one(NF))
+        @test all(progn_new.div[i] .== one(NF))
+        @test all(progn_new.temp[i] .== one(NF))
+        @test all(progn_new.humid[i] .== one(NF))
+        @test all(progn_new.pres[i] .== one(NF))
+    end
+    
+    for (key, value) in progn_new.tracers
+        for value_i in value 
+            @test all(value_i .== one(NF))
+        end 
+    end 
+end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 07ef32c53..aceccd73e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -14,6 +14,7 @@ include("lower_triangular_matrix.jl")
 include("grids.jl")
 include("geodesics.jl")
 include("interpolation.jl")
+include("prognostic_variables.jl")
 include("set.jl")
 
 # GPU/KERNELABSTRACTIONS

From 8026ec5be4e5c048dbcff63e14be1e7de3bb5f0e Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 10 Jan 2025 18:55:46 +0100
Subject: [PATCH 07/62] WIP

---
 test/differentiability/barotropic.jl     | 21 ++++++++++++++++++---
 test/differentiability/timestep_utils.jl | 13 ++++++++-----
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index c17e11519..b1e650c24 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -11,15 +11,30 @@
     progn = prognostic_variables
     diagn = diagnostic_variables
 
-    d_progn = PrognosticVariables(spectral_grid)
-    d_diag = DiagnosticVariables(spectral_grid)
+    diagn_copy = deepcopy(diagn)
+    progn_copy = deepcopy(progn)
+
+    d_progn = zero(progn)
+    d_diag = DiagnosticVariables(spectral_grid, model)
     d_model = deepcopy(model)
 
+    
     #SpeedyWeather.timestep!(progn, diagn, 2Δt, model) # calculate tendencies and leapfrog forward
 
-    autodiff(Reverse, SpeedyWeather.timestep!, Const, Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(2Δt), Duplicated(model, d_model))
+
+    progn_new = zero(progn)
+    dprogn_new = one(progn) # seed 
+
+    autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(2Δt), Duplicated(model, d_model))
 
     # differnetiate wrt initial conditions / previous state
+
+    # new seeed 
+    dprogn_new_2 = one(progn)
+
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, 2Δt, model), dprogn_new_2, progn_copy )
+    
+
     # write this as functions (progn_old, diagn_old, 2\Delta t, model -> progn, diagn)
 
 
diff --git a/test/differentiability/timestep_utils.jl b/test/differentiability/timestep_utils.jl
index 98bae5b0b..ea6ec0d3b 100644
--- a/test/differentiability/timestep_utils.jl
+++ b/test/differentiability/timestep_utils.jl
@@ -1,8 +1,11 @@
-function timestep_ic!(progn_new, diagn_new, progn_old, diagn_old, dt, model)
-
-
+function timestep_oop!(progn_new, progn_old, diagn, dt, model)
+    copy!(progn_new, progn_old)
+    SpeedyWeather.timestep!(progn_new, diagn, dt, model)
+    return nothing
 end 
 
-function timestep_p!() 
-
+function timestep_oop(progn, diagn, dt, model)
+    progn_new = zero(progn)
+    timestep_oop!(progn_new, progn, diagn, dt, model)
+    return progn_new
 end 
\ No newline at end of file

From ed522fa5f5b363eeed6d90b523fbf3990461ead0 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 10 Jan 2025 19:09:04 +0100
Subject: [PATCH 08/62] more VIP step vor

---
 test/differentiability/barotropic.jl     |  3 +++
 test/differentiability/timestep_utils.jl | 18 +++++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index b1e650c24..70c4d375a 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -36,6 +36,9 @@
     
 
     # write this as functions (progn_old, diagn_old, 2\Delta t, model -> progn, diagn)
+    vor = progn.vor[1]
+
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> step_vorticity(x, progn_copy, diagn_copy, 2Δt, model), dprogn_new_2, progn_copy )
 
 
     # differnetiate wrt parameter 
diff --git a/test/differentiability/timestep_utils.jl b/test/differentiability/timestep_utils.jl
index ea6ec0d3b..ccccb7de7 100644
--- a/test/differentiability/timestep_utils.jl
+++ b/test/differentiability/timestep_utils.jl
@@ -8,4 +8,20 @@ function timestep_oop(progn, diagn, dt, model)
     progn_new = zero(progn)
     timestep_oop!(progn_new, progn, diagn, dt, model)
     return progn_new
-end 
\ No newline at end of file
+end 
+
+# for comparisions with FiniteDifferences.jl we need inputs and outputs that are easy to convert
+# to Vectors. PrognosticVariables, and even more so a Model struct are not really easy to convert
+# Vectors. So, for the comparision we look at more limited problems. E.g. a single step and just 
+# looking at the vorticiity
+function step_vorticity!(vorticity_new::LowerTriangularArray, vorticity::LowerTriangularArray, progn, diagn, dt, model)
+    set!(progn, vor=vorticity, lf=1)
+    SpeedyWeather.timestep!(progn, diagn, dt, model)
+    copy!(vorticity_new, progn.vor[1])
+end 
+
+function step_vorticity(vorticity::LowerTriangularArray, progn, diagn, dt, model)
+    vorticity_new = zero(vorticity)
+    step_vorticity!(vorticity_new, vorticity, progn, diagn, dt, model)
+    return vorticity_new
+end
\ No newline at end of file

From 21e150079d2b07e5df5a27cb2d19869a09548b93 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 10 Jan 2025 19:28:40 +0100
Subject: [PATCH 09/62] changelog

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 237bb2b52..5f5aab838 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- zero, fill! and one for PrognosticVariables added [#656](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/656)
 - Differentiability tests for timestepping added [#656](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/656)
 - New output variables definition simplified [#653](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/653)
 - Bug: get_vertices for full grids had southpole at 90˚N [#654](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/654)

From 2db15831c9e36cb3c448f6253b2d90fd4d12a061 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Mon, 13 Jan 2025 13:21:14 +0100
Subject: [PATCH 10/62] more WIP

---
 test/differentiability/barotropic.jl     | 32 +++++++++++++++++++-----
 test/differentiability/timestep_utils.jl |  5 +++-
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index 70c4d375a..b75c1ccf6 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -25,20 +25,40 @@
     progn_new = zero(progn)
     dprogn_new = one(progn) # seed 
 
+    # this is a test of the full timestep of the model. We do this just to see that it causes no errors
+    # and the gradient is non-zero. we don't have something to compare it against
     autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(2Δt), Duplicated(model, d_model))
 
-    # differnetiate wrt initial conditions / previous state
+    @test sum(d_progn.vor[1]) != 0
+    #dprogn_new.vor[1]
+
+    # differnetiate vorticity wrt initial conditions / previous state
 
     # new seeed 
-    dprogn_new_2 = one(progn)
+    vor = progn.vor[1] 
+    vor_copy = deepcopy(vor)
+    dvor = zero(vor)
 
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, 2Δt, model), dprogn_new_2, progn_copy )
-    
+    vor_new = zero(vor)
+    dvor_new = zero(vor)
+    fill!(dvor_new, 1+1im)
+
+    #somethings not right
+
+    diagn_copy = deepcopy(diagn)
+    progn_copy = deepcopy(progn)
+    d_progn = zero(progn)
+    d_diagn = DiagnosticVariables(spectral_grid, model)
+
+    # step_vorticity!(vor_new, vor, progn, diagn, 2Δt, model)
+    autodiff(Reverse, step_vorticity!, Const, Duplicated(vor_new, dvor_new), Duplicated(vor, dvor), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(2Δt), Duplicated(model, d_model))
+
+    # To-DO: there's a problem with copy!(LTA), do something about it 
 
     # write this as functions (progn_old, diagn_old, 2\Delta t, model -> progn, diagn)
-    vor = progn.vor[1]
+    dvor_2 = one(progn.vor[1])
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> step_vorticity(x, progn_copy, diagn_copy, 2Δt, model), dvor_2, vor_copy)
 
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> step_vorticity(x, progn_copy, diagn_copy, 2Δt, model), dprogn_new_2, progn_copy )
 
 
     # differnetiate wrt parameter 
diff --git a/test/differentiability/timestep_utils.jl b/test/differentiability/timestep_utils.jl
index ccccb7de7..886b69438 100644
--- a/test/differentiability/timestep_utils.jl
+++ b/test/differentiability/timestep_utils.jl
@@ -15,9 +15,12 @@ end
 # Vectors. So, for the comparision we look at more limited problems. E.g. a single step and just 
 # looking at the vorticiity
 function step_vorticity!(vorticity_new::LowerTriangularArray, vorticity::LowerTriangularArray, progn, diagn, dt, model)
-    set!(progn, vor=vorticity, lf=1)
+    #set!(progn, model.geometry, vor=vorticity, lf=1)
+    progn.vor[1] .= vorticity
     SpeedyWeather.timestep!(progn, diagn, dt, model)
     copy!(vorticity_new, progn.vor[1])
+    #vorticity_new.data .= progn.vor[1].data
+    return nothing
 end 
 
 function step_vorticity(vorticity::LowerTriangularArray, progn, diagn, dt, model)

From ad43ef729064d1325b63ccc4568da06e8b3e1b46 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Sun, 19 Jan 2025 13:25:10 +0100
Subject: [PATCH 11/62] fix old spectral gradient test

---
 test/differentiability/speedy_transforms.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/differentiability/speedy_transforms.jl b/test/differentiability/speedy_transforms.jl
index 19869e12e..d8f761829 100644
--- a/test/differentiability/speedy_transforms.jl
+++ b/test/differentiability/speedy_transforms.jl
@@ -182,7 +182,7 @@ end
             dv = zero(v)
             div = zero(u)
             ddiv = zero(u)
-            fill!(ddiv, 1)
+            fill!(ddiv, 1 + 1im)
 
             autodiff(Reverse, divergence!, Const, Duplicated(div, ddiv), Duplicated(u, du), Duplicated(v, dv), Duplicated(S, dS))
 
@@ -197,7 +197,7 @@ end
             end 
 
             ddiv2 = zero(ddiv)
-            fill!(ddiv2, 1)
+            fill!(ddiv2, 1 + 1im)
 
             fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> divergence(x[1],x[2], S), ddiv2, (u, v))
             @test isapprox(du, fd_jvp[1][1])

From be8d08ca07bece05c9686f616044b550853e7525 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Sun, 19 Jan 2025 18:11:15 +0100
Subject: [PATCH 12/62] PrognosticVariables compatible with FiniteDifferences

---
 ext/SpeedyWeatherFiniteDifferencesExt.jl | 81 ++++++++++++++++++++++++
 1 file changed, 81 insertions(+)

diff --git a/ext/SpeedyWeatherFiniteDifferencesExt.jl b/ext/SpeedyWeatherFiniteDifferencesExt.jl
index 3527f8737..8f21f13d5 100644
--- a/ext/SpeedyWeatherFiniteDifferencesExt.jl
+++ b/ext/SpeedyWeatherFiniteDifferencesExt.jl
@@ -27,4 +27,85 @@ function FiniteDifferences.to_vec(x::LTA) where LTA <: LowerTriangularArray
     return x_vec, LowerTriangularArray_from_vec
 end
 
+###
+### PrognosticVariables
+###
+
+function flatten(prog::PrognosticVariables{NF, ArrayType, NSTEPS, SpectralVariable2D, SpectralVariable3D, GridVariable2D}) where {NF,ArrayType, NSTEPS, SpectralVariable2D, SpectralVariable3D, GridVariable2D}
+
+    (; trunc, nlayers, nlat_half) = prog
+    nvars = 5 + length(progn.tracers)
+
+    # do it as a LTA 
+    prog_array = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers, NSTEPS, nvars)
+
+    for istep in 1:NSTEPS
+        prog_array[:, :, istep, 1] = prog.vor[istep]
+        prog_array[:, :, istep, 2] = prog.div[istep]
+        prog_array[:, :, istep, 3] = prog.temp[istep]
+        prog_array[:, :, istep, 4] = prog.humid[istep]
+        prog_array[:, 1, istep, 5] = prog.pres[istep]
+
+        for (i_key, (key, values)) in enumerate(prog.tracers)
+            prog_array[:,:, istep, 5+i_key] = values
+        end 
+    end 
+
+    nvars_grid = 6 
+    land_ocean_array = zeros(GridVariable2D, nlat_half, nvars_grid)
+    land_ocean_array[:, 1] = prog.ocean.sea_surface_temperature
+    land_ocean_array[:, 2] = prog.ocean.sea_ice_concentration
+    land_ocean_array[:, 3] = prog.land.land_surface_temperature
+    land_ocean_array[:, 4] = prog.land.snow_depth
+    land_ocean_array[:, 5] = prog.land.soil_moisture_layer1
+    land_ocean_array[:, 6] = prog.land.soil_moisture_layer2
+ 
+    return prog_array, land_ocean_array
+end 
+
+function SpeedyWeather.set!(prog::PrognosticVariables, L::LowerTriangularArray, G::AbstractGridArray)
+    
+    # the geometry isn't actually used for the set! we do 
+    geometry = Geometry(SpectralGrid(trunc = prog.trunc, nlayers=prog.nlayers, Grid=typeof(prog.land.land_surface_temperature)))
+
+    for istep in axes(L, 3)
+        set!(prog, geometry; vor = L[:,:,istep,1], lf=istep)
+        set!(prog, geometry; div = L[:,:,istep,2], lf=istep)
+        set!(prog, geometry; temp = L[:,:,istep,3], lf=istep)
+        set!(prog, geometry; humid = L[:,:,istep,4], lf=istep)
+        set!(prog, geometry; pres = L[:,1,istep,5], lf=istep)
+
+        for (i_key, (key, values)) in enumerate(prog.tracers)
+            values .= L[:,:, istep, 5+i_key] 
+        end 
+    end 
+
+    set!(prog, geometry; sea_surface_temperature = G[:,1])
+    set!(prog, geometry; sea_ice_concentration = G[:,2])
+    set!(prog, geometry; land_surface_temperature = G[:,3])
+    set!(prog, geometry; snow_depth = G[:,4])
+    set!(prog, geometry; soil_moisture_layer1 = G[:,5])
+    set!(prog, geometry; soil_moisture_layer2 = G[:,6])
+
+    return nothing
+end 
+
+function FiniteDifferences.to_vec(prog::PrognosticVariables)
+
+    flattened_prog_lta, flattened_prog_grid = flatten(prog)
+
+    flattened_prog_lta, lta_from_vec = FiniteDifferences.to_vec(flattened_prog_lta)
+    flattened_prog_grid, grid_from_vec = FiniteDifferences.to_vec(flattened_prog_grid)
+
+    N_lta = length(flattened_prog_lta)
+
+    function PrognosticVariables_from_vec(x_vec)
+        prog_new = zero(prog)
+        set!(prog_new, lta_from_vec(x_vec[1:N_lta]), grid_from_vec(x_vec[N_lta+1:end])) 
+        return prog_new 
+    end 
+
+    return cat(flattened_prog_lta, flattened_prog_grid, dims=1), PrognosticVariables_from_vec
+end 
+
 end 
\ No newline at end of file

From 684aef657881aef66c72cb59e5ef609cf80acf73 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Sun, 19 Jan 2025 20:00:32 +0100
Subject: [PATCH 13/62] cleanup timestepping tests a bit, test currenlty
 failing

---
 test/differentiability/barotropic.jl     | 51 ++++++++----------------
 test/differentiability/runtests.jl       | 10 ++++-
 test/differentiability/timestep_utils.jl | 19 ---------
 3 files changed, 25 insertions(+), 55 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index b75c1ccf6..e0afe3935 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -1,12 +1,14 @@
 
 @testset "Differentiability: Barotropic Model Timestepping" begin 
-    spectral_grid = SpectralGrid(trunc=31, nlayers=1)          # define resolution
+    # T15 still yields somewhat sensible dynamics, that's why it's chosen here
+    spectral_grid = SpectralGrid(trunc=15, nlayers=1)          # define resolution
     model = BarotropicModel(; spectral_grid)   # construct model
     simulation = initialize!(model)  
     initialize!(simulation)
 
     (; prognostic_variables, diagnostic_variables, model) = simulation
     (; Δt, Δt_millisec) = model.time_stepping
+    dt = 2Δt
 
     progn = prognostic_variables
     diagn = diagnostic_variables
@@ -18,48 +20,29 @@
     d_diag = DiagnosticVariables(spectral_grid, model)
     d_model = deepcopy(model)
 
-    
-    #SpeedyWeather.timestep!(progn, diagn, 2Δt, model) # calculate tendencies and leapfrog forward
-
-
     progn_new = zero(progn)
     dprogn_new = one(progn) # seed 
 
-    # this is a test of the full timestep of the model. We do this just to see that it causes no errors
-    # and the gradient is non-zero. we don't have something to compare it against
-    autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(2Δt), Duplicated(model, d_model))
-
-    @test sum(d_progn.vor[1]) != 0
-    #dprogn_new.vor[1]
-
-    # differnetiate vorticity wrt initial conditions / previous state
-
-    # new seeed 
-    vor = progn.vor[1] 
-    vor_copy = deepcopy(vor)
-    dvor = zero(vor)
-
-    vor_new = zero(vor)
-    dvor_new = zero(vor)
-    fill!(dvor_new, 1+1im)
+    # test if differentiation works wrt copy! (there were some problems with it before)
+    autodiff(Reverse, copy!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn))
 
-    #somethings not right
+    # it's identity: 
+    @test all(flatten(d_progn)[1] .== 1)
+    @test all(flatten(d_progn)[2] .== 1)
 
-    diagn_copy = deepcopy(diagn)
-    progn_copy = deepcopy(progn)
-    d_progn = zero(progn)
-    d_diagn = DiagnosticVariables(spectral_grid, model)
-
-    # step_vorticity!(vor_new, vor, progn, diagn, 2Δt, model)
-    autodiff(Reverse, step_vorticity!, Const, Duplicated(vor_new, dvor_new), Duplicated(vor, dvor), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(2Δt), Duplicated(model, d_model))
+    progn_new = zero(progn)
+    dprogn_new = one(progn) # seed 
 
-    # To-DO: there's a problem with copy!(LTA), do something about it 
+    # test that we can differentiate wrt an IC 
+    autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
 
-    # write this as functions (progn_old, diagn_old, 2\Delta t, model -> progn, diagn)
-    dvor_2 = one(progn.vor[1])
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> step_vorticity(x, progn_copy, diagn_copy, 2Δt, model), dvor_2, vor_copy)
+    @test sum(to_vec(d_progn)[1]) != 0
 
+    dprogn_2 = one(progn) # seed 
+ 
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
+    @test isapprox(to_vec(fd_jvp[1])[1], isapprox(to_vec(d_progn)[1]))
 
     # differnetiate wrt parameter 
     # write this as function (model, progn, diagn, 2\Delta t) -> progn_new
diff --git a/test/differentiability/runtests.jl b/test/differentiability/runtests.jl
index c0cddb42c..30943077e 100644
--- a/test/differentiability/runtests.jl
+++ b/test/differentiability/runtests.jl
@@ -6,11 +6,17 @@ import EnzymeTestUtils: test_approx
 import FiniteDifferences: j′vp, grad, central_fdm
 import AbstractFFTs
 
+FiniteDiffExt = Base.get_extension(SpeedyWeather, :SpeedyWeatherFiniteDifferencesExt)
+import .FiniteDiffExt: flatten
+
 grid_types = [FullGaussianGrid, OctahedralGaussianGrid] # one full and one reduced grid, both Gaussian to have exact transforms 
 grid_dealiasing = [2, 3]
 fd_tests = [true, true] 
 
-include("speedy_transforms.jl")
-
+# UTILITIES 
+#include("type_utils.jl")
 include("timestep_utils.jl")
+
+# TESTS 
+include("speedy_transforms.jl")
 include("barotropic.jl")
\ No newline at end of file
diff --git a/test/differentiability/timestep_utils.jl b/test/differentiability/timestep_utils.jl
index 886b69438..2fe7cde04 100644
--- a/test/differentiability/timestep_utils.jl
+++ b/test/differentiability/timestep_utils.jl
@@ -9,22 +9,3 @@ function timestep_oop(progn, diagn, dt, model)
     timestep_oop!(progn_new, progn, diagn, dt, model)
     return progn_new
 end 
-
-# for comparisions with FiniteDifferences.jl we need inputs and outputs that are easy to convert
-# to Vectors. PrognosticVariables, and even more so a Model struct are not really easy to convert
-# Vectors. So, for the comparision we look at more limited problems. E.g. a single step and just 
-# looking at the vorticiity
-function step_vorticity!(vorticity_new::LowerTriangularArray, vorticity::LowerTriangularArray, progn, diagn, dt, model)
-    #set!(progn, model.geometry, vor=vorticity, lf=1)
-    progn.vor[1] .= vorticity
-    SpeedyWeather.timestep!(progn, diagn, dt, model)
-    copy!(vorticity_new, progn.vor[1])
-    #vorticity_new.data .= progn.vor[1].data
-    return nothing
-end 
-
-function step_vorticity(vorticity::LowerTriangularArray, progn, diagn, dt, model)
-    vorticity_new = zero(vorticity)
-    step_vorticity!(vorticity_new, vorticity, progn, diagn, dt, model)
-    return vorticity_new
-end
\ No newline at end of file

From c96bd7486eaf267adf77ede9c43c959748ec8368 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Sun, 19 Jan 2025 20:22:23 +0100
Subject: [PATCH 14/62] small fixes and WIP

---
 ext/SpeedyWeatherFiniteDifferencesExt.jl | 2 +-
 test/differentiability/barotropic.jl     | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/ext/SpeedyWeatherFiniteDifferencesExt.jl b/ext/SpeedyWeatherFiniteDifferencesExt.jl
index 8f21f13d5..757c7d537 100644
--- a/ext/SpeedyWeatherFiniteDifferencesExt.jl
+++ b/ext/SpeedyWeatherFiniteDifferencesExt.jl
@@ -34,7 +34,7 @@ end
 function flatten(prog::PrognosticVariables{NF, ArrayType, NSTEPS, SpectralVariable2D, SpectralVariable3D, GridVariable2D}) where {NF,ArrayType, NSTEPS, SpectralVariable2D, SpectralVariable3D, GridVariable2D}
 
     (; trunc, nlayers, nlat_half) = prog
-    nvars = 5 + length(progn.tracers)
+    nvars = 5 + length(prog.tracers)
 
     # do it as a LTA 
     prog_array = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers, NSTEPS, nvars)
diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index e0afe3935..c19350399 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -36,13 +36,15 @@
     # test that we can differentiate wrt an IC 
     autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
 
+    # nonzero gradient
     @test sum(to_vec(d_progn)[1]) != 0
 
+    # FD comparison 
     dprogn_2 = one(progn) # seed 
- 
+
     fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
-    @test isapprox(to_vec(fd_jvp[1])[1], isapprox(to_vec(d_progn)[1]))
+    @test isapprox(to_vec(fd_jvp[1])[1], to_vec(d_progn)[1])
 
     # differnetiate wrt parameter 
     # write this as function (model, progn, diagn, 2\Delta t) -> progn_new

From 9701b28429495fe99a97349fd3009bc0a7ce6282 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Mon, 20 Jan 2025 19:38:47 +0100
Subject: [PATCH 15/62] make_zero works for model

---
 ext/SpeedyWeatherEnzymeExt.jl | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/ext/SpeedyWeatherEnzymeExt.jl b/ext/SpeedyWeatherEnzymeExt.jl
index 5ab34499f..2ddb3c96b 100644
--- a/ext/SpeedyWeatherEnzymeExt.jl
+++ b/ext/SpeedyWeatherEnzymeExt.jl
@@ -2,6 +2,8 @@ module SpeedyWeatherEnzymeExt
 
 using SpeedyWeather
 using Enzyme
+using Enzyme.EnzymeCore
+using SpeedyWeather.ProgressMeter
 import .EnzymeRules: reverse, augmented_primal
 using .EnzymeRules
 
@@ -112,4 +114,17 @@ function reverse(config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(_four
     return (nothing, nothing, nothing, nothing)
 end
 
+###
+# implement make_zero where the default one fails
+
+# this lock is part of the ProgressMeter that's part of the Feedback of all models
+@inline function Enzyme.make_zero(
+    ::Type{ProgressMeter.ProgressCore}, 
+    seen::IdDict, 
+    prev::ProgressMeter.ProgressCore, 
+    ::Val{copy_if_inactive} = Val(false),
+)::ProgressMeter.ProgressCore where {copy_if_inactive} 
+    return prev
+end
+
 end
\ No newline at end of file

From 055f496ef0ef70a96b131be27f60ed467c9eb2a7 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 22 Jan 2025 14:12:51 +0100
Subject: [PATCH 16/62] updated finitedifferences compat

---
 ext/SpeedyWeatherFiniteDifferencesExt.jl | 95 ++++--------------------
 src/dynamics/particles.jl                |  3 +
 2 files changed, 19 insertions(+), 79 deletions(-)

diff --git a/ext/SpeedyWeatherFiniteDifferencesExt.jl b/ext/SpeedyWeatherFiniteDifferencesExt.jl
index 757c7d537..875fe8e2f 100644
--- a/ext/SpeedyWeatherFiniteDifferencesExt.jl
+++ b/ext/SpeedyWeatherFiniteDifferencesExt.jl
@@ -27,85 +27,22 @@ function FiniteDifferences.to_vec(x::LTA) where LTA <: LowerTriangularArray
     return x_vec, LowerTriangularArray_from_vec
 end
 
-###
-### PrognosticVariables
-###
-
-function flatten(prog::PrognosticVariables{NF, ArrayType, NSTEPS, SpectralVariable2D, SpectralVariable3D, GridVariable2D}) where {NF,ArrayType, NSTEPS, SpectralVariable2D, SpectralVariable3D, GridVariable2D}
-
-    (; trunc, nlayers, nlat_half) = prog
-    nvars = 5 + length(prog.tracers)
-
-    # do it as a LTA 
-    prog_array = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers, NSTEPS, nvars)
-
-    for istep in 1:NSTEPS
-        prog_array[:, :, istep, 1] = prog.vor[istep]
-        prog_array[:, :, istep, 2] = prog.div[istep]
-        prog_array[:, :, istep, 3] = prog.temp[istep]
-        prog_array[:, :, istep, 4] = prog.humid[istep]
-        prog_array[:, 1, istep, 5] = prog.pres[istep]
-
-        for (i_key, (key, values)) in enumerate(prog.tracers)
-            prog_array[:,:, istep, 5+i_key] = values
-        end 
-    end 
-
-    nvars_grid = 6 
-    land_ocean_array = zeros(GridVariable2D, nlat_half, nvars_grid)
-    land_ocean_array[:, 1] = prog.ocean.sea_surface_temperature
-    land_ocean_array[:, 2] = prog.ocean.sea_ice_concentration
-    land_ocean_array[:, 3] = prog.land.land_surface_temperature
-    land_ocean_array[:, 4] = prog.land.snow_depth
-    land_ocean_array[:, 5] = prog.land.soil_moisture_layer1
-    land_ocean_array[:, 6] = prog.land.soil_moisture_layer2
- 
-    return prog_array, land_ocean_array
-end 
-
-function SpeedyWeather.set!(prog::PrognosticVariables, L::LowerTriangularArray, G::AbstractGridArray)
-    
-    # the geometry isn't actually used for the set! we do 
-    geometry = Geometry(SpectralGrid(trunc = prog.trunc, nlayers=prog.nlayers, Grid=typeof(prog.land.land_surface_temperature)))
-
-    for istep in axes(L, 3)
-        set!(prog, geometry; vor = L[:,:,istep,1], lf=istep)
-        set!(prog, geometry; div = L[:,:,istep,2], lf=istep)
-        set!(prog, geometry; temp = L[:,:,istep,3], lf=istep)
-        set!(prog, geometry; humid = L[:,:,istep,4], lf=istep)
-        set!(prog, geometry; pres = L[:,1,istep,5], lf=istep)
-
-        for (i_key, (key, values)) in enumerate(prog.tracers)
-            values .= L[:,:, istep, 5+i_key] 
-        end 
+# needs an extra modification because an empty vector yields Any[] with to_vec for Particle (which isn't the case for number types)
+function FiniteDifferences.to_vec(x::Vector{Particle{NF}}) where NF 
+    if isempty(x) 
+        return NF[], identity
+    else # the else statement is the unmodified to_vec(::DenseVector)
+        x_vecs_and_backs = map(to_vec, x)
+        x_vecs, backs = first.(x_vecs_and_backs), last.(x_vecs_and_backs)
+        function Vector_from_vec(x_vec)
+            sz = cumsum(map(length, x_vecs))
+            x_Vec = [backs[n](x_vec[sz[n] - length(x_vecs[n]) + 1:sz[n]]) for n in eachindex(x)]
+            return oftype(x, x_Vec)
+        end
+        # handle empty x
+        x_vec = isempty(x_vecs) ? eltype(eltype(x_vecs))[] : reduce(vcat, x_vecs)
+        return x_vec, Vector_from_vec
     end 
-
-    set!(prog, geometry; sea_surface_temperature = G[:,1])
-    set!(prog, geometry; sea_ice_concentration = G[:,2])
-    set!(prog, geometry; land_surface_temperature = G[:,3])
-    set!(prog, geometry; snow_depth = G[:,4])
-    set!(prog, geometry; soil_moisture_layer1 = G[:,5])
-    set!(prog, geometry; soil_moisture_layer2 = G[:,6])
-
-    return nothing
-end 
-
-function FiniteDifferences.to_vec(prog::PrognosticVariables)
-
-    flattened_prog_lta, flattened_prog_grid = flatten(prog)
-
-    flattened_prog_lta, lta_from_vec = FiniteDifferences.to_vec(flattened_prog_lta)
-    flattened_prog_grid, grid_from_vec = FiniteDifferences.to_vec(flattened_prog_grid)
-
-    N_lta = length(flattened_prog_lta)
-
-    function PrognosticVariables_from_vec(x_vec)
-        prog_new = zero(prog)
-        set!(prog_new, lta_from_vec(x_vec[1:N_lta]), grid_from_vec(x_vec[N_lta+1:end])) 
-        return prog_new 
-    end 
-
-    return cat(flattened_prog_lta, flattened_prog_grid, dims=1), PrognosticVariables_from_vec
-end 
+end
 
 end 
\ No newline at end of file
diff --git a/src/dynamics/particles.jl b/src/dynamics/particles.jl
index 4a932dff7..fce92089a 100644
--- a/src/dynamics/particles.jl
+++ b/src/dynamics/particles.jl
@@ -48,6 +48,9 @@ function Base.zeros(ArrayType::Type{<:AbstractArray{P}}, n::Int...) where {P<:Pa
     fill!(z, zero(P))
 end
 
+Base.eltype(::Type{Particle{NF}}) where NF = NF
+Base.eltype(::Particle{NF}) where NF = NF
+
 Base.rand(rng::Random.AbstractRNG, ::Random.Sampler{Particle}) = rand(rng, Particle{DEFAULT_NF,true})
 Base.rand(rng::Random.AbstractRNG, ::Random.Sampler{Particle{NF}}) where NF = rand(rng, Particle{NF,true})
 

From e0d44acb6dc77fd665c6f45743eff6bdda07f40b Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 22 Jan 2025 17:19:41 +0100
Subject: [PATCH 17/62] FiniteDifferences now working with Diagn and Progn

---
 ext/SpeedyWeatherFiniteDifferencesExt.jl | 59 +++++++++++++++++++++++-
 test/differentiability/runtests.jl       |  7 +--
 test/differentiability/test_utils.jl     | 32 +++++++++++++
 3 files changed, 92 insertions(+), 6 deletions(-)
 create mode 100644 test/differentiability/test_utils.jl

diff --git a/ext/SpeedyWeatherFiniteDifferencesExt.jl b/ext/SpeedyWeatherFiniteDifferencesExt.jl
index 875fe8e2f..97840c251 100644
--- a/ext/SpeedyWeatherFiniteDifferencesExt.jl
+++ b/ext/SpeedyWeatherFiniteDifferencesExt.jl
@@ -27,7 +27,7 @@ function FiniteDifferences.to_vec(x::LTA) where LTA <: LowerTriangularArray
     return x_vec, LowerTriangularArray_from_vec
 end
 
-# needs an extra modification because an empty vector yields Any[] with to_vec for Particle (which isn't the case for number types)
+# Vector{Particle} needs an extra modification because an empty vector yields Any[] with to_vec for Particle (which isn't the case for number types)
 function FiniteDifferences.to_vec(x::Vector{Particle{NF}}) where NF 
     if isempty(x) 
         return NF[], identity
@@ -45,4 +45,61 @@ function FiniteDifferences.to_vec(x::Vector{Particle{NF}}) where NF
     end 
 end
 
+# A version of the generic fallback from FiniteDifferences that excludes some of the fields 
+# that we don't want to be varied for our big data structures 
+function FiniteDifferences.to_vec(x::T) where {T <: Union{PrognosticVariables, PrognosticVariablesOcean, PrognosticVariablesLand, DiagnosticVariables, Tendencies, GridVariables, DynamicsVariables, PhysicsVariables, ParticleVariables}}
+
+    val_vecs_and_backs = map(name -> to_vec(getfield(x, name)), included_fields(T))
+    vals = first.(val_vecs_and_backs)
+    backs = last.(val_vecs_and_backs)
+
+    vals_excluded_pre = map(name -> getfield(x, name), excluded_fields_pre(T))
+    vals_excluded_post = map(name -> getfield(x, name), excluded_fields_post(T))
+
+    v, vals_from_vec = to_vec(vals)
+    function structtype_from_vec(v::Vector{<:Real})
+        val_vecs = vals_from_vec(v)
+        values = map((b, v) -> b(v), backs, val_vecs)
+        
+        T(vals_excluded_pre..., values..., vals_excluded_post...)
+    end
+    return v, structtype_from_vec
+end
+
+included_fields(::Type{<:PrognosticVariables}) = (:vor, :div, :temp, :humid, :pres, :random_pattern, :ocean, :land, :tracers, :particles)
+excluded_fields_pre(::Type{<:PrognosticVariables}) = (:trunc, :nlat_half, :nlayers, :nparticles)
+excluded_fields_post(::Type{<:PrognosticVariables}) = (:scale, :clock)
+
+included_fields(::Type{<:PrognosticVariablesOcean}) = (:sea_surface_temperature, :sea_ice_concentration)
+excluded_fields_pre(::Type{<:PrognosticVariablesOcean}) = (:nlat_half, )
+excluded_fields_post(::Type{<:PrognosticVariablesOcean}) = ()
+
+included_fields(::Type{<:PrognosticVariablesLand}) = (:land_surface_temperature, :snow_depth, :soil_moisture_layer1, :soil_moisture_layer2)
+excluded_fields_pre(::Type{<:PrognosticVariablesLand}) = (:nlat_half, )
+excluded_fields_post(::Type{<:PrognosticVariablesLand}) = ()
+
+included_fields(::Type{<:DiagnosticVariables}) = (:tendencies, :grid, :dynamics, :physics, :column, :temp_average)
+excluded_fields_pre(::Type{<:DiagnosticVariables}) = (:trunc, :nlat_half, :nlayers, :nparticle)
+excluded_fields_post(::Type{<:DiagnosticVariables}) = (:scale,)
+
+included_fields(::Type{<:Tendencies}) = (:vor_tend, :div_tend, :temp_tend, :humid_tend, :u_tend, :v_tend, :pres_tend, :tracers_tend, :u_tend_grid, :v_tend_grid, :temp_tend_grid, :humid_tend_grid, :pres_tend_grid, :tracers_tend_grid)
+excluded_fields_pre(::Type{<:Tendencies}) = (:trunc, :nlat_half, :nlayers)
+excluded_fields_post(::Type{<:Tendencies}) = ()
+
+included_fields(::Type{<:GridVariables}) = (:vor_grid, :div_grid, :temp_grid, :temp_virt_grid, :humid_grid, :u_grid, :v_grid, :pres_grid, :tracers_grid, :random_pattern, :temp_grid_prev, :humid_grid_prev, :u_grid_prev, :v_grid_prev, :pres_grid_prev, :tracers_grid_prev)
+excluded_fields_pre(::Type{<:GridVariables}) = (:nlat_half, :nlayers)
+excluded_fields_post(::Type{<:GridVariables}) = ()
+
+included_fields(::Type{<:DynamicsVariables}) = (:a, :b, :a_grid, :b_grid, :a_2D, :b_2D, :a_2D_grid, :b_2D_grid, :uv∇lnp, :uv∇lnp_sum_above, :div_sum_above, :temp_virt, :geopot, :σ_tend, :∇lnp_x, :∇lnp_y, :u_mean_grid, :v_mean_grid, :div_mean_grid, :div_mean)
+excluded_fields_pre(::Type{<:DynamicsVariables}) = (:trunc, :nlat_half, :nlayers)
+excluded_fields_post(::Type{<:DynamicsVariables}) = ()
+
+included_fields(::Type{<:PhysicsVariables}) = (:precip_large_scale, :precip_convection, :cloud_top, :soil_moisture_availability, :surface_flux_heat, :surface_flux_humid, :outgoing_shortwave_radiation, :outgoing_longwave_radiation, :cos_zenith)
+excluded_fields_pre(::Type{<:PhysicsVariables}) = (:nlat_half,)
+excluded_fields_post(::Type{<:PhysicsVariables}) = ()
+
+included_fields(::Type{<:ParticleVariables}) = (:locations, :u, :v, :σ_tend)
+excluded_fields_pre(::Type{<:ParticleVariables}) = (:nparticles, :nlat_half)
+excluded_fields_post(::Type{<:ParticleVariables}) = (:interpolator, )
+
 end 
\ No newline at end of file
diff --git a/test/differentiability/runtests.jl b/test/differentiability/runtests.jl
index 30943077e..fe925ded9 100644
--- a/test/differentiability/runtests.jl
+++ b/test/differentiability/runtests.jl
@@ -1,20 +1,17 @@
 # Here we test the differentiability of the model by comparing Enzyme's results to finite difference differenation 
 # These tests are relatively expensive, and also not strictly necessary to perform at every commit, so they 
 # are only part of the extended test set 
-using SpeedyWeather, EnzymeTestUtils, Enzyme, FiniteDifferences
+using SpeedyWeather, EnzymeTestUtils, Enzyme, FiniteDifferences, Test
 import EnzymeTestUtils: test_approx
 import FiniteDifferences: j′vp, grad, central_fdm
 import AbstractFFTs
 
-FiniteDiffExt = Base.get_extension(SpeedyWeather, :SpeedyWeatherFiniteDifferencesExt)
-import .FiniteDiffExt: flatten
-
 grid_types = [FullGaussianGrid, OctahedralGaussianGrid] # one full and one reduced grid, both Gaussian to have exact transforms 
 grid_dealiasing = [2, 3]
 fd_tests = [true, true] 
 
 # UTILITIES 
-#include("type_utils.jl")
+include("test_utils.jl")
 include("timestep_utils.jl")
 
 # TESTS 
diff --git a/test/differentiability/test_utils.jl b/test/differentiability/test_utils.jl
new file mode 100644
index 000000000..92049a8a8
--- /dev/null
+++ b/test/differentiability/test_utils.jl
@@ -0,0 +1,32 @@
+# this is mainly there for testing purposes, and a bit more manual then above 
+function flatten(prog::PrognosticVariables{NF, ArrayType, NSTEPS, SpectralVariable2D, SpectralVariable3D, GridVariable2D}) where {NF,ArrayType, NSTEPS, SpectralVariable2D, SpectralVariable3D, GridVariable2D}
+
+    (; trunc, nlayers, nlat_half) = prog
+    nvars = 5 + length(progn.tracers)
+
+    # do it as a LTA 
+    prog_array = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers, NSTEPS, nvars)
+
+    for istep in 1:NSTEPS
+        prog_array[:, :, istep, 1] = prog.vor[istep]
+        prog_array[:, :, istep, 2] = prog.div[istep]
+        prog_array[:, :, istep, 3] = prog.temp[istep]
+        prog_array[:, :, istep, 4] = prog.humid[istep]
+        prog_array[:, 1, istep, 5] = prog.pres[istep]
+
+        for (i_key, (key, values)) in enumerate(prog.tracers)
+            prog_array[:,:, istep, 5+i_key] = values
+        end 
+    end 
+
+    nvars_grid = 6 
+    land_ocean_array = zeros(GridVariable2D, nlat_half, nvars_grid)
+    land_ocean_array[:, 1] = prog.ocean.sea_surface_temperature
+    land_ocean_array[:, 2] = prog.ocean.sea_ice_concentration
+    land_ocean_array[:, 3] = prog.land.land_surface_temperature
+    land_ocean_array[:, 4] = prog.land.snow_depth
+    land_ocean_array[:, 5] = prog.land.soil_moisture_layer1
+    land_ocean_array[:, 6] = prog.land.soil_moisture_layer2
+
+    return prog_array, land_ocean_array
+end 

From 8ff6b3a9d09e5740ac557fa4894d6424082fa112 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 22 Jan 2025 18:31:14 +0100
Subject: [PATCH 18/62] minor confusion about finitedifference correctness

---
 test/differentiability/barotropic.jl | 65 +++++++++++++++++++++++++++-
 1 file changed, 63 insertions(+), 2 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index c19350399..528aff954 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -17,8 +17,8 @@
     progn_copy = deepcopy(progn)
 
     d_progn = zero(progn)
-    d_diag = DiagnosticVariables(spectral_grid, model)
-    d_model = deepcopy(model)
+    d_diag = make_zero(diagn)
+    d_model = make_zero(model)
 
     progn_new = zero(progn)
     dprogn_new = one(progn) # seed 
@@ -46,6 +46,67 @@
 
     @test isapprox(to_vec(fd_jvp[1])[1], to_vec(d_progn)[1])
 
+    lf1 = 2 
+    lf2 = 2 
+
+    # set the tendencies back to zero for accumulation
+    fill!(diagn.tendencies, 0, Barotropic)
+
+    # TENDENCIES, DIFFUSION, LEAPFROGGING AND TRANSFORM SPECTRAL STATE TO GRID
+    SpeedyWeather.dynamics_tendencies!(diagn, progn, lf2, model)
+    SpeedyWeather.horizontal_diffusion!(diagn, progn, model.horizontal_diffusion, model)
+
+ 
+    progn_copy = deepcopy(progn)
+    dprogn = one(progn_copy)
+    dprogn_copy = one(progn_copy)
+
+    tend = diagn.tendencies
+    tend_copy = deepcopy(tend)
+    dtend = make_zero(tend)
+    dmodel = make_zero(model)
+
+    #leapfrog!(progn, diagn.tendencies, dt, lf1, model)
+    autodiff(Reverse, SpeedyWeather.leapfrog!, Const, Duplicated(progn, dprogn), Duplicated(tend, dtend), Const(dt), Const(lf1), Const(model))
+
+    function leapfrog(progn, tend, dt, lf, model)
+        SpeedyWeather.leapfrog!(progn, tend, dt, lf, model)
+        return progn
+    end 
+
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> leapfrog(progn, x, dt, lf1, model), dprogn_copy, tend_copy)
+
+    # single variable leapfrog step 
+
+    A_old = progn.vor[1]
+    A_old_copy = copy(A_old)
+    dA_old = one(A_old)
+
+    A_new = progn.vor[2]
+    A_new_copy = copy(A_new)
+    dA_new = one(A_new)
+
+    tendency = diagn.tendencies.vor_tend
+    tendency_copy = copy(tendency)
+    dtendency = zero(tendency)
+
+    L = model.time_stepping
+
+    autodiff(Reverse, SpeedyWeather.leapfrog!, Const, Duplicated(A_old, dA_old), Duplicated(A_new, dA_new), Duplicated(tendency, dtendency), Const(dt), Const(lf1), Const(L))
+
+    # d tend needs to be: dt* ( 1 + w1 - w2) (for every coefficient)
+    # enzyme shows it is 
+    
+    function leapfrog_single(A_old, A_new, tendency, dt, lf, L)
+
+        A_old_new = copy(A_old)
+        A_new_new = copy(A_new)
+        SpeedyWeather.leapfrog!(A_old_new, A_new_new, tendency, dt, lf, L)
+        return A_new_new
+    end 
+
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> leapfrog_single(A_old_copy, A_new_copy, x, dt, lf1, L), one(A_new_copy), tendency_copy)
+
     # differnetiate wrt parameter 
     # write this as function (model, progn, diagn, 2\Delta t) -> progn_new
 

From 5e414aa055678643f3e870a86b99a2a59d2a15e6 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Thu, 23 Jan 2025 18:26:15 +0100
Subject: [PATCH 19/62] working towards timestepping diff PrimitiveModel

---
 src/dynamics/implicit.jl               |  8 +--
 src/physics/surface_fluxes.jl          | 19 +++---
 src/physics/vertical_diffusion.jl      |  8 +--
 src/physics/zenith.jl                  |  4 +-
 test/differentiability/timestepping.jl | 93 ++++++++++++++++++++++++++
 5 files changed, 113 insertions(+), 19 deletions(-)
 create mode 100644 test/differentiability/timestepping.jl

diff --git a/src/dynamics/implicit.jl b/src/dynamics/implicit.jl
index 694fb9f23..b470b4581 100644
--- a/src/dynamics/implicit.jl
+++ b/src/dynamics/implicit.jl
@@ -209,12 +209,12 @@ Initialize the implicit terms for the PrimitiveEquation models."""
 function initialize!(   
     implicit::ImplicitPrimitiveEquation,
     dt::Real,                                           # the scaled time step radius*dt
-    diagn::DiagnosticVariables,
+    diagn::DiagnosticVariables{NF},
     geometry::AbstractGeometry,
     geopotential::AbstractGeopotential,
     atmosphere::AbstractAtmosphere,
     adiabatic_conversion::AbstractAdiabaticConversion,
-)
+) where NF
 
     (; trunc, nlayers, α, temp_profile, S, S⁻¹, L, R, U, W, L0, L1, L2, L3, L4) = implicit
     (; σ_levels_full, σ_levels_thick) = geometry
@@ -267,10 +267,10 @@ function initialize!(
 
         for r in 1:nlayers
             L1[k, r] = ΔT_below*σ_levels_thick[r]*σₖ         # vert advection operator below
-            L1[k, r] -= k>=r ? σ_levels_thick[r] : 0
+            L1[k, r] -= k>=r ? σ_levels_thick[r] : zero(NF)
 
             L1[k, r] += ΔT_above*σ_levels_thick[r]*σₖ_above   # vert advection operator above
-            L1[k, r] -= (k-1)>=r ? σ_levels_thick[r] : 0
+            L1[k, r] -= (k-1)>=r ? σ_levels_thick[r] : zero(NF)
         end
 
         # _sum_above operator itself
diff --git a/src/physics/surface_fluxes.jl b/src/physics/surface_fluxes.jl
index ba053455d..21be1d4ed 100644
--- a/src/physics/surface_fluxes.jl
+++ b/src/physics/surface_fluxes.jl
@@ -132,10 +132,10 @@ SurfaceHeatFlux(SG::SpectralGrid; kwargs...) = SurfaceHeatFlux{SG.NF}(; kwargs..
 initialize!(::SurfaceHeatFlux, ::PrimitiveEquation) = nothing
 
 function surface_heat_flux!(   
-    column::ColumnVariables,
+    column::ColumnVariables{NF},
     heat_flux::SurfaceHeatFlux,
     model::PrimitiveEquation,
-)   
+) where NF
     cₚ = model.atmosphere.heat_capacity
     (; heat_exchange_land, heat_exchange_sea) = heat_flux
 
@@ -147,21 +147,22 @@ function surface_heat_flux!(
     land_fraction = column.land_fraction
 
     # drag coefficient
+    # the convert might seem redundant but without it Enzyme struggles with its type and activity 
     drag_sea, drag_land = heat_flux.use_boundary_layer_drag ?
                         (column.boundary_layer_drag, column.boundary_layer_drag) : 
                         (heat_exchange_sea, heat_exchange_land)
 
     # SPEEDY documentation Eq. 54 and 56, land/sea fraction included
     flux_land = ρ*drag_land*V₀*cₚ*(T_skin_land - T)*land_fraction
-    flux_sea  = ρ*drag_sea*V₀*cₚ*(T_skin_sea  - T)*(1-land_fraction)
+    flux_sea  = ρ*drag_sea*V₀*cₚ*(T_skin_sea  - T)*(one(NF)-land_fraction)
 
     # mix fluxes for fractional land-sea mask
     land_available = isfinite(T_skin_land)
     sea_available = isfinite(T_skin_sea)
 
     # Only flux from land/sea if available (not NaN) otherwise zero flux
-    column.flux_temp_upward[end] += land_available ? flux_land : 0
-    column.flux_temp_upward[end] += sea_available ? flux_sea : 0
+    column.flux_temp_upward[end] += land_available ? flux_land : zero(NF)
+    column.flux_temp_upward[end] += sea_available ? flux_sea : zero(NF)
 
     return nothing
 end
@@ -193,9 +194,9 @@ end
 SurfaceEvaporation(SG::SpectralGrid; kwargs...) = SurfaceEvaporation{SG.NF}(; kwargs...)
 initialize!(::SurfaceEvaporation, ::PrimitiveWet) = nothing
 
-function surface_evaporation!(  column::ColumnVariables,
+function surface_evaporation!(  column::ColumnVariables{NF},
                                 evaporation::SurfaceEvaporation,
-                                model::PrimitiveWet)
+                                model::PrimitiveWet) where NF
 
     (; skin_temperature_sea, skin_temperature_land, pres) = column
     (; moisture_exchange_land, moisture_exchange_sea) = evaporation
@@ -226,8 +227,8 @@ function surface_evaporation!(  column::ColumnVariables,
 
     # Only flux from land/sea if available (not NaN) otherwise zero flux
     # mix fluxes for fractional land-sea mask
-    column.flux_humid_upward[end] += land_available ? flux_land : 0
-    column.flux_humid_upward[end] += sea_available ? flux_sea : 0
+    column.flux_humid_upward[end] += land_available ? flux_land : zero(NF)
+    column.flux_humid_upward[end] += sea_available ? flux_sea : zero(NF)
 
     return nothing
 end
\ No newline at end of file
diff --git a/src/physics/vertical_diffusion.jl b/src/physics/vertical_diffusion.jl
index 2c36c9284..ceae9c032 100644
--- a/src/physics/vertical_diffusion.jl
+++ b/src/physics/vertical_diffusion.jl
@@ -110,11 +110,11 @@ function vertical_diffusion!(
 end
 
 function get_diffusion_coefficients!(
-    column::ColumnVariables,
+    column::ColumnVariables{NF},
     scheme::BulkRichardsonDiffusion,
     atmosphere::AbstractAtmosphere,
     planet::AbstractPlanet,
-)
+) where NF
     K = column.b        # reuse work array for diffusion coefficients
     (; Ri_c, κ, z₀, fb) = scheme
     logZ_z₀ = scheme.logZ_z₀[]
@@ -148,11 +148,11 @@ function get_diffusion_coefficients!(
             K_k = K0 * zmin             # = κ*u_N*√Cz in eq. (19, 20)
 
             # multiply with z-dependent factor in eq. (18) ?
-            K_k *= z < fb*h ? 1 : zfac(z, h, fb)
+            K_k *= z < fb*h ? one(NF) : zfac(z, h, fb)
 
             # multiply with Ri-dependent factor in eq. (20) ?
             # TODO use Ri[kₕ] or Ri_N here? 
-            K_k *= Ri[kₕ] <= 0 ? 1 : Rifac(Ri[kₕ], Ri_c, logZ_z₀)
+            K_k *= Ri[kₕ] <= 0 ? one(NF) : Rifac(Ri[kₕ], Ri_c, logZ_z₀)
             K[k] = K_k                  # write diffusion coefficient into array
         end
     else
diff --git a/src/physics/zenith.jl b/src/physics/zenith.jl
index 54d0299a8..ed5b96bed 100644
--- a/src/physics/zenith.jl
+++ b/src/physics/zenith.jl
@@ -211,7 +211,7 @@ function cos_zenith!(
     g = year_angle(NF, time_of_year, length_of_day, length_of_year)
 
     # time correction [radians] due to the equation of time (sunrise/set oscillation)
-    tc = S.equation_of_time ? S.time_correction(g) : 0
+    tc = S.equation_of_time ? S.time_correction(g) : zero(NF)
     
     # solar hour angle at 0˚E (longtiude offset added later)
     λ = 0
@@ -284,7 +284,7 @@ function cos_zenith!(
         ϕ = lat[j]
         h₀ = abs(δ) + abs(ϕ) < π/2 ?        # polar day/night?
         acos(-tan(ϕ) * tan(δ)) :            # if not: calculate length of day
-        ϕ*δ > 0 ? π : 0                     # polar day if signs are equal, otherwise polar night
+        ϕ*δ > 0 ? π : zero(NF)              # polar day if signs are equal, otherwise polar night
         
         sinϕ, cosϕ = sinlat[j], coslat[j]
         cos_zenith_j = h₀*sinδ*sinϕ + cosδ*cosϕ*sin(h₀)
diff --git a/test/differentiability/timestepping.jl b/test/differentiability/timestepping.jl
new file mode 100644
index 000000000..a268473b5
--- /dev/null
+++ b/test/differentiability/timestepping.jl
@@ -0,0 +1,93 @@
+# High-level tests whether time stepping in all models work
+
+@testset "Differentiability: Timestepping" begin 
+    # T15 still yields somewhat sensible dynamics, that's why it's chosen here
+    
+    model_types = [ShallowWaterModel, PrimitiveDryModel, PrimitiveWetModel]
+    model_type = PrimitiveDryModel
+    for model_type in model_types 
+        
+        nlayer = model_type == ShallowWaterModel ? 1 : 3 
+
+        spectral_grid = SpectralGrid(trunc=15, nlayers=nlayer)          # define resolution
+        model = PrimitiveDryModel(; spectral_grid)   # construct model
+        simulation = initialize!(model)  
+        initialize!(simulation)
+
+        (; prognostic_variables, diagnostic_variables, model) = simulation
+        (; Δt, Δt_millisec) = model.time_stepping
+        dt = 2Δt
+
+        progn = prognostic_variables
+        diagn = diagnostic_variables
+
+        diagn_copy = deepcopy(diagn)
+        progn_copy = deepcopy(progn)
+
+        d_progn = zero(progn)
+        d_diag = make_zero(diagn)
+        d_model = make_zero(model)
+
+        progn_new = zero(progn)
+        dprogn_new = one(progn) # seed 
+
+        # test that we can differentiate wrt an IC 
+        autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
+
+        # nonzero gradient
+        @test sum(to_vec(d_progn)[1]) != 0
+
+        # FD comparison 
+        dprogn_2 = one(progn) # seed 
+
+        fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
+
+        @test isapprox(to_vec(fd_jvp[1])[1], to_vec(d_progn)[1])
+
+
+        column = diagn.column 
+        dcolumn = make_zero(column)
+        heat_flux = model.surface_heat_flux 
+        
+        autodiff(Reverse, surface_heat_flux!, Const, Duplicated(column, dcolumn), Const(heat_flux), Const(model))
+
+    end 
+
+end
+
+
+function surface_heat_flux!(   
+    column::ColumnVariables{NF},
+    heat_flux::SurfaceHeatFlux,
+    model::PrimitiveEquation,
+) where NF
+    cₚ = model.atmosphere.heat_capacity
+    (; heat_exchange_land, heat_exchange_sea) = heat_flux
+
+    ρ = column.surface_air_density
+    V₀ = column.surface_wind_speed
+    T_skin_sea = column.skin_temperature_sea
+    T_skin_land = column.skin_temperature_land
+    T = column.surface_temp
+    land_fraction = column.land_fraction
+
+    # drag coefficient
+    # the convert might seem redundant but without it Enzyme struggles with its type and activity 
+    drag_sea, drag_land = heat_flux.use_boundary_layer_drag ?
+                        (column.boundary_layer_drag, column.boundary_layer_drag) : 
+                        (heat_exchange_sea, heat_exchange_land)
+
+    # SPEEDY documentation Eq. 54 and 56, land/sea fraction included
+    flux_land = ρ*drag_land*V₀*cₚ*(T_skin_land - T)*land_fraction
+    flux_sea  = ρ*drag_sea*V₀*cₚ*(T_skin_sea  - T)*(one(NF)-land_fraction)
+
+    # mix fluxes for fractional land-sea mask
+    land_available = isfinite(T_skin_land)
+    sea_available = isfinite(T_skin_sea)
+
+    # Only flux from land/sea if available (not NaN) otherwise zero flux
+    column.flux_temp_upward[end] += land_available ? flux_land : zero(NF)
+    column.flux_temp_upward[end] += sea_available ? flux_sea : zero(NF)
+
+    return nothing
+end
\ No newline at end of file

From 28372aec54f67242154bf52a77e9fdfd12796561 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 24 Jan 2025 11:41:21 +0100
Subject: [PATCH 20/62] minor test update

---
 test/differentiability/barotropic.jl   | 26 +++++++++----
 test/differentiability/runtests.jl     |  3 +-
 test/differentiability/timestepping.jl | 52 ++------------------------
 3 files changed, 23 insertions(+), 58 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index 528aff954..8520c6e90 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -1,4 +1,4 @@
-
+### Experiments going a bit deeper into the timestepping of the barotropic model
 @testset "Differentiability: Barotropic Model Timestepping" begin 
     # T15 still yields somewhat sensible dynamics, that's why it's chosen here
     spectral_grid = SpectralGrid(trunc=15, nlayers=1)          # define resolution
@@ -46,6 +46,8 @@
 
     @test isapprox(to_vec(fd_jvp[1])[1], to_vec(d_progn)[1])
 
+    ### Test the leapfrog 
+
     lf1 = 2 
     lf2 = 2 
 
@@ -56,7 +58,7 @@
     SpeedyWeather.dynamics_tendencies!(diagn, progn, lf2, model)
     SpeedyWeather.horizontal_diffusion!(diagn, progn, model.horizontal_diffusion, model)
 
- 
+
     progn_copy = deepcopy(progn)
     dprogn = one(progn_copy)
     dprogn_copy = one(progn_copy)
@@ -69,12 +71,15 @@
     #leapfrog!(progn, diagn.tendencies, dt, lf1, model)
     autodiff(Reverse, SpeedyWeather.leapfrog!, Const, Duplicated(progn, dprogn), Duplicated(tend, dtend), Const(dt), Const(lf1), Const(model))
 
-    function leapfrog(progn, tend, dt, lf, model)
+    function leapfrog_step(progn::PrognosticVariables, tend, dt, lf, model)
         SpeedyWeather.leapfrog!(progn, tend, dt, lf, model)
         return progn
     end 
 
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> leapfrog(progn, x, dt, lf1, model), dprogn_copy, tend_copy)
+    tend_one = deepcopy(tend_copy)
+    fill!(tend_one, 1, Barotropic)
+
+    fd_jvp = FiniteDifferences.j'vp(central_fdm(5,1), x -> leapfrog_step(progn, x, dt, lf1, model),(tend_copy, tend_one))
 
     # single variable leapfrog step 
 
@@ -94,20 +99,25 @@
 
     autodiff(Reverse, SpeedyWeather.leapfrog!, Const, Duplicated(A_old, dA_old), Duplicated(A_new, dA_new), Duplicated(tendency, dtendency), Const(dt), Const(lf1), Const(L))
 
+    w1 = L.robert_filter*L.williams_filter/2   
+    w2 = L.robert_filter*(1-L.williams_filter)/2   
+    @test all(dtendency .== dt*(1+w1-w2))
     # d tend needs to be: dt* ( 1 + w1 - w2) (for every coefficient)
     # enzyme shows it is 
     
-    function leapfrog_single(A_old, A_new, tendency, dt, lf, L)
-
+    function leapfrog(A_old, A_new, tendency, dt, lf, L)
         A_old_new = copy(A_old)
         A_new_new = copy(A_new)
         SpeedyWeather.leapfrog!(A_old_new, A_new_new, tendency, dt, lf, L)
         return A_new_new
     end 
 
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> leapfrog_single(A_old_copy, A_new_copy, x, dt, lf1, L), one(A_new_copy), tendency_copy)
+    fd_jvp = FiniteDifferences.jvp(central_fdm(5,1), (x) ->leapfrog(A_old_copy, A_old_copy, x, dt, lf1, L), (tendency_copy, one(A_new_copy)))
+    # in this case it need to be dt*(1 - w2) (no contributation from A_old in FD)
 
+    
     # differnetiate wrt parameter 
     # write this as function (model, progn, diagn, 2\Delta t) -> progn_new
+end
+
 
-end
\ No newline at end of file
diff --git a/test/differentiability/runtests.jl b/test/differentiability/runtests.jl
index fe925ded9..992a41251 100644
--- a/test/differentiability/runtests.jl
+++ b/test/differentiability/runtests.jl
@@ -16,4 +16,5 @@ include("timestep_utils.jl")
 
 # TESTS 
 include("speedy_transforms.jl")
-include("barotropic.jl")
\ No newline at end of file
+include("barotropic.jl")
+include("timestepping.jl")
\ No newline at end of file
diff --git a/test/differentiability/timestepping.jl b/test/differentiability/timestepping.jl
index a268473b5..4da7cc9bd 100644
--- a/test/differentiability/timestepping.jl
+++ b/test/differentiability/timestepping.jl
@@ -2,15 +2,14 @@
 
 @testset "Differentiability: Timestepping" begin 
     # T15 still yields somewhat sensible dynamics, that's why it's chosen here
-    
     model_types = [ShallowWaterModel, PrimitiveDryModel, PrimitiveWetModel]
-    model_type = PrimitiveDryModel
+
     for model_type in model_types 
         
-        nlayer = model_type == ShallowWaterModel ? 1 : 3 
+        nlayer = model_type == ShallowWaterModel ? 1 : 3
 
         spectral_grid = SpectralGrid(trunc=15, nlayers=nlayer)          # define resolution
-        model = PrimitiveDryModel(; spectral_grid)   # construct model
+        model = model_type(; spectral_grid)   # construct model
         simulation = initialize!(model)  
         initialize!(simulation)
 
@@ -43,51 +42,6 @@
         fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
         @test isapprox(to_vec(fd_jvp[1])[1], to_vec(d_progn)[1])
-
-
-        column = diagn.column 
-        dcolumn = make_zero(column)
-        heat_flux = model.surface_heat_flux 
-        
-        autodiff(Reverse, surface_heat_flux!, Const, Duplicated(column, dcolumn), Const(heat_flux), Const(model))
-
     end 
 
-end
-
-
-function surface_heat_flux!(   
-    column::ColumnVariables{NF},
-    heat_flux::SurfaceHeatFlux,
-    model::PrimitiveEquation,
-) where NF
-    cₚ = model.atmosphere.heat_capacity
-    (; heat_exchange_land, heat_exchange_sea) = heat_flux
-
-    ρ = column.surface_air_density
-    V₀ = column.surface_wind_speed
-    T_skin_sea = column.skin_temperature_sea
-    T_skin_land = column.skin_temperature_land
-    T = column.surface_temp
-    land_fraction = column.land_fraction
-
-    # drag coefficient
-    # the convert might seem redundant but without it Enzyme struggles with its type and activity 
-    drag_sea, drag_land = heat_flux.use_boundary_layer_drag ?
-                        (column.boundary_layer_drag, column.boundary_layer_drag) : 
-                        (heat_exchange_sea, heat_exchange_land)
-
-    # SPEEDY documentation Eq. 54 and 56, land/sea fraction included
-    flux_land = ρ*drag_land*V₀*cₚ*(T_skin_land - T)*land_fraction
-    flux_sea  = ρ*drag_sea*V₀*cₚ*(T_skin_sea  - T)*(one(NF)-land_fraction)
-
-    # mix fluxes for fractional land-sea mask
-    land_available = isfinite(T_skin_land)
-    sea_available = isfinite(T_skin_sea)
-
-    # Only flux from land/sea if available (not NaN) otherwise zero flux
-    column.flux_temp_upward[end] += land_available ? flux_land : zero(NF)
-    column.flux_temp_upward[end] += sea_available ? flux_sea : zero(NF)
-
-    return nothing
 end
\ No newline at end of file

From c6561c587a93cf89d037d7c3d94456bb45b37dfd Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Sat, 25 Jan 2025 15:52:06 +0100
Subject: [PATCH 21/62] copy! for Clock

---
 ext/SpeedyWeatherEnzymeExt.jl        |  4 ++--
 src/dynamics/clock.jl                | 12 ++++++++++++
 src/dynamics/prognostic_variables.jl |  2 +-
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/ext/SpeedyWeatherEnzymeExt.jl b/ext/SpeedyWeatherEnzymeExt.jl
index 2ddb3c96b..ad19a0e48 100644
--- a/ext/SpeedyWeatherEnzymeExt.jl
+++ b/ext/SpeedyWeatherEnzymeExt.jl
@@ -33,9 +33,9 @@ end
 # Computes the scale for the adjoint/pullback of a real discrete fourier transform.
 function rfft_adjoint_scale(n_freq::Int, n_real::Int)
     if iseven(n_real)
-        return [1; [2 for i=2:(n_freq-1)]; 1]
+        return [1 < i < n_freq ? 2 : 1 for i=1:n_freq]
     else 
-        return [1; [2 for i=2:n_freq]]
+        return [1 < i ? 2 : 1 for i=1:n_freq]
     end 
 end 
 
diff --git a/src/dynamics/clock.jl b/src/dynamics/clock.jl
index 4aa1f1834..188e976b7 100644
--- a/src/dynamics/clock.jl
+++ b/src/dynamics/clock.jl
@@ -38,6 +38,18 @@ function Base.show(io::IO, C::Clock)
     print_fields(io, C, keys)
 end
 
+# copy! 
+function Base.copy!(clock::Clock, clock_old::Clock)
+    clock.time = clock_old.time 
+    clock.start = clock_old.start 
+    clock.period = clock_old.period 
+    clock.timestep_counter = clock_old.timestep_counter
+    clock.n_timesteps = clock_old.n_timesteps
+    clock.Δt = clock_old.Δt 
+    
+    return nothing 
+end 
+
 """
 $(TYPEDSIGNATURES)
 Initialize the clock with the time step `Δt` in the `time_stepping`."""
diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl
index 6556bbc42..0f2f99915 100644
--- a/src/dynamics/prognostic_variables.jl
+++ b/src/dynamics/prognostic_variables.jl
@@ -211,7 +211,7 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
 
     progn_new.random_pattern .= progn_old.random_pattern
 
-    progn_new.clock.time = progn_old.clock.time
+    copy!(progn_new.clock, progn_old.clock)
     progn_new.scale[] = progn_old.scale[]
 
     return progn_new

From 5158af75cc5eb771395dfedded055346d38656b8 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Sat, 25 Jan 2025 18:27:38 +0100
Subject: [PATCH 22/62] avoid ringgrids broadcasting

---
 ext/SpeedyWeatherEnzymeExt.jl            | 10 +++++-----
 src/dynamics/tendencies.jl               | 18 +++++++++---------
 test/differentiability/timestep_utils.jl |  2 +-
 3 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/ext/SpeedyWeatherEnzymeExt.jl b/ext/SpeedyWeatherEnzymeExt.jl
index ad19a0e48..99417a2ee 100644
--- a/ext/SpeedyWeatherEnzymeExt.jl
+++ b/ext/SpeedyWeatherEnzymeExt.jl
@@ -19,15 +19,15 @@ function adjoint_scale(S::SpectralTransform)
     (; nlat_half, nlons, rfft_plans) = S
     nfreqs = [rfft_plan.osz[1] for rfft_plan in rfft_plans] # TODO: This works with FFTW, but does it with cuFFT as well?
 
-    scale = zeros(Int, maximum(nfreqs), nlat_half) 
+    scale = zeros(Int, maximum(nfreqs), 1, nlat_half) # the scratch memory is (Freq x lvl x lat), so we insert 
+                                                      # an additional dimension here for easier matrix multiply
 
     for i=1:nlat_half
-        scale[1:nfreqs[i],i] = rfft_adjoint_scale(nfreqs[i], nlons[i])
+        scale[1:nfreqs[i],1,i] = rfft_adjoint_scale(nfreqs[i], nlons[i])
     end 
 
-    # TODO: transfer array to GPU in case we are on GPU
-    return reshape(scale, maximum(nfreqs), 1, nlat_half) # the scratch memory is (Freq x lvl x lat), so we insert 
-                                                         # an additional dimension here for easier matrix multiply
+    # TODO: transfer array to GPU in case we are on GPU?
+    return scale
 end 
 
 # Computes the scale for the adjoint/pullback of a real discrete fourier transform.
diff --git a/src/dynamics/tendencies.jl b/src/dynamics/tendencies.jl
index c7785ace1..888e75553 100644
--- a/src/dynamics/tendencies.jl
+++ b/src/dynamics/tendencies.jl
@@ -838,8 +838,8 @@ function SpeedyTransforms.transform!(
     model::PrimitiveEquation;
     initialize::Bool = false,
 )   
-    (; vor_grid, div_grid, pres_grid, u_grid, v_grid, temp_grid, humid_grid) = diagn.grid
-    (; temp_grid_prev, humid_grid_prev, u_grid_prev, v_grid_prev) = diagn.grid
+    (; vor_grid, div_grid, pres_grid, u_grid, v_grid, temp_grid, humid_grid, 
+    temp_grid_prev, humid_grid_prev, u_grid_prev, v_grid_prev) = diagn.grid
 
     vor   = progn.vor[lf]           # relative vorticity at leapfrog step lf
     div   = progn.div[lf]           # divergence at leapfrog step lf
@@ -853,10 +853,10 @@ function SpeedyTransforms.transform!(
 
     # retain previous time step for vertical advection and some parameterizations
     if initialize == false              # only store prev after initial step
-        @. temp_grid_prev = temp_grid   # this is temperature anomaly wrt to implicit reference profile!
-        @. humid_grid_prev = humid_grid
-        @. u_grid_prev = u_grid
-        @. v_grid_prev = v_grid
+        @. temp_grid_prev.data = temp_grid.data   # this is temperature anomaly wrt to implicit reference profile!
+        @. humid_grid_prev.data = humid_grid.data
+        @. u_grid_prev.data = u_grid.data
+        @. v_grid_prev.data = v_grid.data
 
         for (name, tracer) in model.tracers
             if tracer.active
@@ -900,9 +900,9 @@ function SpeedyTransforms.transform!(
             end
         end
 
-        @. humid_grid_prev = humid_grid
-        @. u_grid_prev = u_grid
-        @. v_grid_prev = v_grid
+        @. humid_grid_prev.data = humid_grid.data
+        @. u_grid_prev.data = u_grid.data
+        @. v_grid_prev.data = v_grid.data
 
         for (name, tracer) in model.tracers
             if tracer.active
diff --git a/test/differentiability/timestep_utils.jl b/test/differentiability/timestep_utils.jl
index 2fe7cde04..0986058a7 100644
--- a/test/differentiability/timestep_utils.jl
+++ b/test/differentiability/timestep_utils.jl
@@ -1,4 +1,4 @@
-function timestep_oop!(progn_new, progn_old, diagn, dt, model)
+function timestep_oop!(progn_new::PrognosticVariables, progn_old::PrognosticVariables, diagn, dt, model)
     copy!(progn_new, progn_old)
     SpeedyWeather.timestep!(progn_new, diagn, dt, model)
     return nothing

From 92e42fe316af759f17a973baedbd2bc8e7c19ecb Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Mon, 27 Jan 2025 15:03:17 +0100
Subject: [PATCH 23/62] include one differentiability test in main test set

---
 src/RingGrids/general.jl            |  5 ++--
 src/dynamics/tendencies.jl          |  1 +
 test/spectral_transform_ad_rules.jl | 37 +++++++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl
index 026e457f3..1c1dd975f 100644
--- a/src/RingGrids/general.jl
+++ b/src/RingGrids/general.jl
@@ -427,7 +427,6 @@ end
 ## BROADCASTING
 # following https://docs.julialang.org/en/v1/manual/interfaces/#man-interfaces-broadcasting
 import Base.Broadcast: BroadcastStyle, Broadcasted, DefaultArrayStyle
-import LinearAlgebra: isstructurepreserving, fzeropreserving
 
 # {1} as grids are <:AbstractVector, Grid here is the non-parameteric Grid type!
 struct AbstractGridArrayStyle{N, Grid} <: Broadcast.AbstractArrayStyle{N} end
@@ -440,7 +439,7 @@ Base.BroadcastStyle(::Type{Grid}) where {Grid<:AbstractGridArray{T, N, ArrayType
 
 # allocation for broadcasting, create a new Grid with undef of type/number format T
 function Base.similar(bc::Broadcasted{AbstractGridArrayStyle{N, Grid}}, ::Type{T}) where {N, Grid, T}
-    return Grid(Array{T}(undef, size(bc)...))
+    return Grid(Array{T}(undef, size(bc)))
 end
 
 # ::Val{0} for broadcasting with 0-dimensional, ::Val{1} for broadcasting with vectors, etc
@@ -485,7 +484,7 @@ function Base.similar(
     ::Type{T},
 ) where {N, ArrayType, Grid, T}
     ArrayType_ = nonparametric_type(ArrayType)
-    return Grid(ArrayType_{T}(undef, size(bc)...))
+    return Grid(ArrayType_{T}(undef, size(bc)))
 end
 
 function Adapt.adapt_structure(to, grid::Grid) where {Grid <: AbstractGridArray}
diff --git a/src/dynamics/tendencies.jl b/src/dynamics/tendencies.jl
index 888e75553..4ee807b34 100644
--- a/src/dynamics/tendencies.jl
+++ b/src/dynamics/tendencies.jl
@@ -857,6 +857,7 @@ function SpeedyTransforms.transform!(
         @. humid_grid_prev.data = humid_grid.data
         @. u_grid_prev.data = u_grid.data
         @. v_grid_prev.data = v_grid.data
+        # TODO: above broadcastings are done with the .data because there's currently a problem with Enzyme and our RingGrids broadcasting
 
         for (name, tracer) in model.tracers
             if tracer.active
diff --git a/test/spectral_transform_ad_rules.jl b/test/spectral_transform_ad_rules.jl
index 3a68c0062..9bf393c68 100644
--- a/test/spectral_transform_ad_rules.jl
+++ b/test/spectral_transform_ad_rules.jl
@@ -56,4 +56,41 @@ end
     @testset "Complete Transform ChainRules" begin 
         # WIP
     end
+end 
+
+@testset "Complete Differentiability" begin 
+    # We do extensive correctness checks and tests on the differentiability 
+    # in a seperate test set. But we do want to ensure in the regular CI that 
+    # we don't commit some kind of problem for the Enzyme differentiability
+    # so, we test here if we get a non-zero gradient from the timestepping.  
+    function timestep_oop!(progn_new::PrognosticVariables, progn_old::PrognosticVariables, diagn, dt, model)
+        copy!(progn_new, progn_old)
+        SpeedyWeather.timestep!(progn_new, diagn, dt, model)
+        return nothing
+    end 
+
+    spectral_grid = SpectralGrid(trunc=15, nlayers=3)          # define resolution
+    model = PrimitiveWetModel(; spectral_grid)   # construct model
+    simulation = initialize!(model)  
+    initialize!(simulation)
+
+    (; prognostic_variables, diagnostic_variables, model) = simulation
+    (; Δt, Δt_millisec) = model.time_stepping
+    dt = 2Δt
+
+    progn = prognostic_variables
+    diagn = diagnostic_variables
+
+    diagn_copy = deepcopy(diagn)
+    progn_copy = deepcopy(progn)
+
+    d_progn = zero(progn)
+    d_diag = make_zero(diagn)
+    d_model = make_zero(model)
+
+    progn_new = zero(progn)
+    dprogn_new = one(progn) # seed 
+
+    autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
+    @test sum(to_vec(d_progn)[1]) != 0
 end 
\ No newline at end of file

From 52aad60c32495fa2f2938e9de27bcbf72bfa1585 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Mon, 27 Jan 2025 18:38:33 +0100
Subject: [PATCH 24/62] leapfrog correctness passing

---
 test/differentiability/barotropic.jl | 34 ++++++++++++++++------------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index 8520c6e90..af20365d1 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -46,7 +46,9 @@
 
     @test isapprox(to_vec(fd_jvp[1])[1], to_vec(d_progn)[1])
 
-    ### Test the leapfrog 
+    #
+    # Test the leapfrog 
+    # 
 
     lf1 = 2 
     lf2 = 2 
@@ -58,7 +60,6 @@
     SpeedyWeather.dynamics_tendencies!(diagn, progn, lf2, model)
     SpeedyWeather.horizontal_diffusion!(diagn, progn, model.horizontal_diffusion, model)
 
-
     progn_copy = deepcopy(progn)
     dprogn = one(progn_copy)
     dprogn_copy = one(progn_copy)
@@ -68,20 +69,23 @@
     dtend = make_zero(tend)
     dmodel = make_zero(model)
 
-    #leapfrog!(progn, diagn.tendencies, dt, lf1, model)
     autodiff(Reverse, SpeedyWeather.leapfrog!, Const, Duplicated(progn, dprogn), Duplicated(tend, dtend), Const(dt), Const(lf1), Const(model))
 
-    function leapfrog_step(progn::PrognosticVariables, tend, dt, lf, model)
-        SpeedyWeather.leapfrog!(progn, tend, dt, lf, model)
-        return progn
+    function leapfrog_step(progn_new::PrognosticVariables, progn::PrognosticVariables, tend, dt, lf, model)
+        copy!(progn_new, progn)
+        SpeedyWeather.leapfrog!(progn_new, tend, dt, lf, model)
+        return progn_new
     end 
 
-    tend_one = deepcopy(tend_copy)
-    fill!(tend_one, 1, Barotropic)
+    prog_new = zero(progn_copy)
+
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> leapfrog_step(prog_new, progn_copy, x, dt, lf1, model), dprogn_copy, tend_copy)
 
-    fd_jvp = FiniteDifferences.j'vp(central_fdm(5,1), x -> leapfrog_step(progn, x, dt, lf1, model),(tend_copy, tend_one))
+    @test all(isapprox.(to_vec(fd_jvp[1])[1], to_vec(dtend)[1],rtol=1e-2,atol=1e-2))
 
+    # 
     # single variable leapfrog step 
+    # 
 
     A_old = progn.vor[1]
     A_old_copy = copy(A_old)
@@ -101,10 +105,9 @@
 
     w1 = L.robert_filter*L.williams_filter/2   
     w2 = L.robert_filter*(1-L.williams_filter)/2   
-    @test all(dtendency .== dt*(1+w1-w2))
+    @test all(dtendency .≈ dt*(1+w1-w2))
     # d tend needs to be: dt* ( 1 + w1 - w2) (for every coefficient)
-    # enzyme shows it is 
-    
+
     function leapfrog(A_old, A_new, tendency, dt, lf, L)
         A_old_new = copy(A_old)
         A_new_new = copy(A_new)
@@ -112,10 +115,13 @@
         return A_new_new
     end 
 
-    fd_jvp = FiniteDifferences.jvp(central_fdm(5,1), (x) ->leapfrog(A_old_copy, A_old_copy, x, dt, lf1, L), (tendency_copy, one(A_new_copy)))
+    #fd_jvp = FiniteDifferences.jvp(central_fdm(5,1), (x) ->leapfrog(A_old_copy, A_new_copy, x, dt, lf1, L), (tendency_copy, one(tendency_copy)))
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), (x) ->leapfrog(A_old_copy, A_new_copy, x, dt, lf1, L), one(tendency_copy), tendency_copy)
+    
     # in this case it need to be dt*(1 - w2) (no contributation from A_old in FD)
+    @test all(isapprox.(fd_jvp[1], dt*(1-w2), rtol=1e-2))
+
 
-    
     # differnetiate wrt parameter 
     # write this as function (model, progn, diagn, 2\Delta t) -> progn_new
 end

From 4d5b0ef610719ef5fd35c05dd1e74dbec5e61159 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Mon, 27 Jan 2025 19:39:21 +0100
Subject: [PATCH 25/62] fix FD ext, towards transfrom diag correctnes test

---
 ext/SpeedyWeatherFiniteDifferencesExt.jl |  6 ++---
 test/differentiability/barotropic.jl     | 28 +++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/ext/SpeedyWeatherFiniteDifferencesExt.jl b/ext/SpeedyWeatherFiniteDifferencesExt.jl
index 97840c251..6a5b5e6aa 100644
--- a/ext/SpeedyWeatherFiniteDifferencesExt.jl
+++ b/ext/SpeedyWeatherFiniteDifferencesExt.jl
@@ -78,8 +78,8 @@ included_fields(::Type{<:PrognosticVariablesLand}) = (:land_surface_temperature,
 excluded_fields_pre(::Type{<:PrognosticVariablesLand}) = (:nlat_half, )
 excluded_fields_post(::Type{<:PrognosticVariablesLand}) = ()
 
-included_fields(::Type{<:DiagnosticVariables}) = (:tendencies, :grid, :dynamics, :physics, :column, :temp_average)
-excluded_fields_pre(::Type{<:DiagnosticVariables}) = (:trunc, :nlat_half, :nlayers, :nparticle)
+included_fields(::Type{<:DiagnosticVariables}) = (:tendencies, :grid, :dynamics, :physics, :particles, :column, :temp_average)
+excluded_fields_pre(::Type{<:DiagnosticVariables}) = (:trunc, :nlat_half, :nlayers, :nparticles)
 excluded_fields_post(::Type{<:DiagnosticVariables}) = (:scale,)
 
 included_fields(::Type{<:Tendencies}) = (:vor_tend, :div_tend, :temp_tend, :humid_tend, :u_tend, :v_tend, :pres_tend, :tracers_tend, :u_tend_grid, :v_tend_grid, :temp_tend_grid, :humid_tend_grid, :pres_tend_grid, :tracers_tend_grid)
@@ -94,7 +94,7 @@ included_fields(::Type{<:DynamicsVariables}) = (:a, :b, :a_grid, :b_grid, :a_2D,
 excluded_fields_pre(::Type{<:DynamicsVariables}) = (:trunc, :nlat_half, :nlayers)
 excluded_fields_post(::Type{<:DynamicsVariables}) = ()
 
-included_fields(::Type{<:PhysicsVariables}) = (:precip_large_scale, :precip_convection, :cloud_top, :soil_moisture_availability, :surface_flux_heat, :surface_flux_humid, :outgoing_shortwave_radiation, :outgoing_longwave_radiation, :cos_zenith)
+included_fields(::Type{<:PhysicsVariables}) = (:precip_large_scale, :precip_convection, :precip_rate_large_scale, :precip_rate_convection, :cloud_top, :soil_moisture_availability, :sensible_heat_flux, :evaporative_flux, :surface_shortwave_up, :surface_shortwave_down, :surface_longwave_up, :surface_longwave_down, :outgoing_shortwave_radiation, :outgoing_longwave_radiation, :cos_zenith)
 excluded_fields_pre(::Type{<:PhysicsVariables}) = (:nlat_half,)
 excluded_fields_post(::Type{<:PhysicsVariables}) = ()
 
diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index af20365d1..62f5101d7 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -116,12 +116,38 @@
     end 
 
     #fd_jvp = FiniteDifferences.jvp(central_fdm(5,1), (x) ->leapfrog(A_old_copy, A_new_copy, x, dt, lf1, L), (tendency_copy, one(tendency_copy)))
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), (x) ->leapfrog(A_old_copy, A_new_copy, x, dt, lf1, L), one(tendency_copy), tendency_copy)
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x ->leapfrog(A_old_copy, A_new_copy, x, dt, lf1, L), one(tendency_copy), tendency_copy)
     
     # in this case it need to be dt*(1 - w2) (no contributation from A_old in FD)
     @test all(isapprox.(fd_jvp[1], dt*(1-w2), rtol=1e-2))
 
+    #
+    # transform!(diagn, progn, lf2, model)
+    #
+
+    fill!(diagn.tendencies, 0, PrimitiveWetModel)
+    diag_copy = deepcopy(diagn)
+    
+    ddiag = deepcopy(diagn)
+    fill!(ddiag.tendencies, 1, PrimitiveWetModel)
+
+    # full one, also including grid data, temp_average, column 
+    ddiag_copy = deepcopy(ddiag)
+
+    progn_copy = deepcopy(progn)
+    dprogn = make_zero(progn)
+
+    autodiff(Reverse, SpeedyWeather.transform!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(lf2), Duplicated(model, make_zero(model)))
+
+    function transform_diagn(diag, progn, lf2, model)
+        diag_copy = deepcopy(diag)
+        transform!(diag_copy, progn, lf2, model)
+        return diag_copy
+    end 
+
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_diagn(diag_copy, x, lf2, model), ddiag_copy, progn_copy)
 
+ 
     # differnetiate wrt parameter 
     # write this as function (model, progn, diagn, 2\Delta t) -> progn_new
 end

From ea96381eb16fb4dc11a60a34cb82b878c3f890de Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Tue, 28 Jan 2025 10:26:44 +0100
Subject: [PATCH 26/62] first doc sketch

---
 docs/make.jl                  | 1 +
 docs/src/differentiability.md | 8 ++++++++
 2 files changed, 9 insertions(+)
 create mode 100644 docs/src/differentiability.md

diff --git a/docs/make.jl b/docs/make.jl
index d6fe1dafb..4e26a1fea 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -20,6 +20,7 @@ makedocs(
                 "Stochastic physics" => "stochastic_physics.md",
                 "Analysis"=>"analysis.md",
                 "Tree structure"=>"structure.md",
+                "Differentiability and Adjoint Model"=>"differentiability.md",
                 "NetCDF output"=>"output.md",
             ],
             "Extending SpeedyWeather" => [
diff --git a/docs/src/differentiability.md b/docs/src/differentiability.md
new file mode 100644
index 000000000..3008608cc
--- /dev/null
+++ b/docs/src/differentiability.md
@@ -0,0 +1,8 @@
+# Differentiability and Adjoint Model
+
+SpeedyWeather.jl is written with differentiability in mind. This means that our model is differentiable by automatic differentiation (AD). If you are interested in machine learning (ML), this means that you can integrate our model directly into your ML models without the need to train your ANNs seperately first. For atmospheric modellers this means that you get an adjoint model for free which is always generated automatically, so that we don't need to maintain it seperatly. So, you can calibrate SpeedyWeather.jl in a fully automatic, data-driven way. 
+
+!!! Work in progress
+    The differentiability of SpeedyWeather.jl is still work in progress and some parts of this documentation might be not be always updated to the latest state. We will extend this documentation over time.  
+
+For the differentiability of our model we rely on [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl). If you've used Enzyme before, just go ahead and try to differnentiate the model! It should work. We have checked the correctness of the gradients extensively against a finite differences differentiation with [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl/). In the following we present a simple example how we can take the gradient of a single timestep of the primitive equation model with respect to one of the model parameter, so with ``\mathbf{P}_t = (\zeta, \mathcal{D}, p_s, T, q)`` as the state vector of all prognostic variables at time step
\ No newline at end of file

From 60e28fae532241bd94072045b306dc420e9afeb4 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Tue, 28 Jan 2025 10:27:07 +0100
Subject: [PATCH 27/62] more correctness tests

---
 test/differentiability/barotropic.jl | 53 +++++++++++++++++++++++-----
 1 file changed, 45 insertions(+), 8 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index 62f5101d7..dcb32ecd2 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -45,6 +45,48 @@
     fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
     @test isapprox(to_vec(fd_jvp[1])[1], to_vec(d_progn)[1])
+    
+    #
+    # We go individually through all components of the time stepping and check 
+    # correctness
+    #
+
+    fill!(diagn.tendencies, 0, Barotropic)
+
+    #
+    # dynamics_tendencies!
+    #
+
+    lf2 = 
+    diagn_copy = deepcopy(diagn)
+    ddiag = one(diagn_copy)
+    ddiag_copy = deepcopy(ddiag)
+    progn_copy = deepcopy(progn)
+
+    #autodiff(Reverse, SpeedyWeather.dynamics_tendencies!(diagn, progn, lf2, model)
+
+    #
+    # horizontal_diffusion!
+    #
+    lf1 = 1
+    diagn_copy = deepcopy(diagn)
+    ddiag = one(diagn_copy)
+    ddiag_copy = deepcopy(ddiag)
+
+    progn_copy = deepcopy(progn)
+    dprogn = one(progn)
+
+    autodiff(Reverse, SpeedyWeather.horizontal_diffusion!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(model.horizontal_diffusion), Duplicated(model, make_zero(model)), Const(lf1))
+
+    function horizontal_diffusion(diagn, progn, diffusion, model, lf)
+        diagn_new = deepcopy(diagn)
+        SpeedyWeather.horizontal_diffusion!(diagn_new, progn, diffusion, model, lf)
+        return diagn_new
+    end 
+
+    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> horizontal_diffusion(diag_copy, x, model.horizontal_diffusion, model, lf1), ddiag_copy, progn_copy)
+    # look closer here! 
+    @test all(isapprox.(to_vec(fd_jvp[1])[1], to_vec(dprogn)[1],rtol=1e-2,atol=1e-2))
 
     #
     # Test the leapfrog 
@@ -54,10 +96,8 @@
     lf2 = 2 
 
     # set the tendencies back to zero for accumulation
-    fill!(diagn.tendencies, 0, Barotropic)
 
     # TENDENCIES, DIFFUSION, LEAPFROGGING AND TRANSFORM SPECTRAL STATE TO GRID
-    SpeedyWeather.dynamics_tendencies!(diagn, progn, lf2, model)
     SpeedyWeather.horizontal_diffusion!(diagn, progn, model.horizontal_diffusion, model)
 
     progn_copy = deepcopy(progn)
@@ -115,7 +155,6 @@
         return A_new_new
     end 
 
-    #fd_jvp = FiniteDifferences.jvp(central_fdm(5,1), (x) ->leapfrog(A_old_copy, A_new_copy, x, dt, lf1, L), (tendency_copy, one(tendency_copy)))
     fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x ->leapfrog(A_old_copy, A_new_copy, x, dt, lf1, L), one(tendency_copy), tendency_copy)
     
     # in this case it need to be dt*(1 - w2) (no contributation from A_old in FD)
@@ -128,10 +167,7 @@
     fill!(diagn.tendencies, 0, PrimitiveWetModel)
     diag_copy = deepcopy(diagn)
     
-    ddiag = deepcopy(diagn)
-    fill!(ddiag.tendencies, 1, PrimitiveWetModel)
-
-    # full one, also including grid data, temp_average, column 
+    ddiag = one(diagn)
     ddiag_copy = deepcopy(ddiag)
 
     progn_copy = deepcopy(progn)
@@ -146,8 +182,9 @@
     end 
 
     fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_diagn(diag_copy, x, lf2, model), ddiag_copy, progn_copy)
+    
+    @test all(isapprox.(to_vec(fd_jvp[1])[1], to_vec(dprogn)[1],rtol=1e-2,atol=1e-2))
 
- 
     # differnetiate wrt parameter 
     # write this as function (model, progn, diagn, 2\Delta t) -> progn_new
 end

From 141dc28af2d7f446205ae93021e38e71a18800ae Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Tue, 28 Jan 2025 10:27:27 +0100
Subject: [PATCH 28/62] one for diag based on to_vec

---
 test/differentiability/test_utils.jl | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/test/differentiability/test_utils.jl b/test/differentiability/test_utils.jl
index 92049a8a8..fe44cbba1 100644
--- a/test/differentiability/test_utils.jl
+++ b/test/differentiability/test_utils.jl
@@ -30,3 +30,9 @@ function flatten(prog::PrognosticVariables{NF, ArrayType, NSTEPS, SpectralVariab
 
     return prog_array, land_ocean_array
 end 
+
+function Base.one(diag::DiagnosticVariables{NF}) where NF
+    vec, re = to_vec(diag)
+    vec .= NF(1)
+    return re(vec)
+end 
\ No newline at end of file

From 6a9f6cca7831d24fd650649b030ef80baa8b2449 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Tue, 28 Jan 2025 16:04:53 +0100
Subject: [PATCH 29/62] finish correctness test individual functions barotropic
 model

---
 docs/src/differentiability.md               |  2 +-
 test/differentiability/barotropic.jl        | 60 +++++++++++++++------
 test/differentiability/speedy_transforms.jl | 42 +++++++--------
 3 files changed, 67 insertions(+), 37 deletions(-)

diff --git a/docs/src/differentiability.md b/docs/src/differentiability.md
index 3008608cc..ad55fc03d 100644
--- a/docs/src/differentiability.md
+++ b/docs/src/differentiability.md
@@ -3,6 +3,6 @@
 SpeedyWeather.jl is written with differentiability in mind. This means that our model is differentiable by automatic differentiation (AD). If you are interested in machine learning (ML), this means that you can integrate our model directly into your ML models without the need to train your ANNs seperately first. For atmospheric modellers this means that you get an adjoint model for free which is always generated automatically, so that we don't need to maintain it seperatly. So, you can calibrate SpeedyWeather.jl in a fully automatic, data-driven way. 
 
 !!! Work in progress
-    The differentiability of SpeedyWeather.jl is still work in progress and some parts of this documentation might be not be always updated to the latest state. We will extend this documentation over time.  
+    The differentiability of SpeedyWeather.jl is still work in progress and some parts of this documentation might be not be always updated to the latest state. We will extend this documentation over time. Don't hesitate to contact us via GitHub issues or mail when you have questions. 
 
 For the differentiability of our model we rely on [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl). If you've used Enzyme before, just go ahead and try to differnentiate the model! It should work. We have checked the correctness of the gradients extensively against a finite differences differentiation with [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl/). In the following we present a simple example how we can take the gradient of a single timestep of the primitive equation model with respect to one of the model parameter, so with ``\mathbf{P}_t = (\zeta, \mathcal{D}, p_s, T, q)`` as the state vector of all prognostic variables at time step
\ No newline at end of file
diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index dcb32ecd2..1b13b3f9d 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -42,9 +42,9 @@
     # FD comparison 
     dprogn_2 = one(progn) # seed 
 
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
-    @test isapprox(to_vec(fd_jvp[1])[1], to_vec(d_progn)[1])
+    @test isapprox(to_vec(fd_vjp[1])[1], to_vec(d_progn)[1])
     
     #
     # We go individually through all components of the time stepping and check 
@@ -57,24 +57,40 @@
     # dynamics_tendencies!
     #
 
-    lf2 = 
+    lf2 = 2 
+
     diagn_copy = deepcopy(diagn)
     ddiag = one(diagn_copy)
     ddiag_copy = deepcopy(ddiag)
     progn_copy = deepcopy(progn)
+    dprogn = make_zero(progn)
+
+    autodiff(Reverse, SpeedyWeather.dynamics_tendencies!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(lf2), Duplicated(model, make_zero(model)))
+
+    function dynamics_tendencies(diagn, progn, lf, model)
+        diagn_new = deepcopy(diagn)
+        SpeedyWeather.dynamics_tendencies!(diagn_new, progn, lf, model)
+        return diagn_new
+    end 
 
-    #autodiff(Reverse, SpeedyWeather.dynamics_tendencies!(diagn, progn, lf2, model)
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> dynamics_tendencies(diagn_copy, x, lf2, model), ddiag_copy, progn_copy)
+
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-2))
+
+    # in the default configuration without forcing or drag, the barotropic model's don't dependent on the previous prognostic state 
+    @test sum(to_vec(dprogn)[1]) ≈ 0 
 
     #
     # horizontal_diffusion!
     #
+
     lf1 = 1
     diagn_copy = deepcopy(diagn)
     ddiag = one(diagn_copy)
     ddiag_copy = deepcopy(ddiag)
 
     progn_copy = deepcopy(progn)
-    dprogn = one(progn)
+    dprogn = make_zero(progn)
 
     autodiff(Reverse, SpeedyWeather.horizontal_diffusion!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(model.horizontal_diffusion), Duplicated(model, make_zero(model)), Const(lf1))
 
@@ -84,9 +100,23 @@
         return diagn_new
     end 
 
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> horizontal_diffusion(diag_copy, x, model.horizontal_diffusion, model, lf1), ddiag_copy, progn_copy)
-    # look closer here! 
-    @test all(isapprox.(to_vec(fd_jvp[1])[1], to_vec(dprogn)[1],rtol=1e-2,atol=1e-2))
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> horizontal_diffusion(diagn_copy, x, model.horizontal_diffusion, model, lf1), ddiag_copy, progn_copy)
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-2))
+
+    # ∂(progn)
+    # should be row-wise `model.horizontal_diffusion.impl .* model.horizontal_diffusion.expl`
+    # for all variables that are diffused 
+    diff_coefficient = model.horizontal_diffusion.impl .* model.horizontal_diffusion.expl
+    l_indices = [(1:l) for l=1:progn.vor[1].n]
+    for (i,il) in enumerate(l_indices)
+        @test all(real.(Matrix(dprogn.vor[lf1][:,1])[i, il]) .≈ diff_coefficient[i])
+    end 
+
+    # ∂(tend_old)
+    # should be row-wise `model.horizontal_diffusion.impl` 
+    for (i,il) in enumerate(l_indices)
+        @test all(real.(Matrix(ddiag.tendencies.vor_tend[:,1])[i, il]) .≈ model.horizontal_diffusion.impl[i])
+    end 
 
     #
     # Test the leapfrog 
@@ -119,9 +149,9 @@
 
     prog_new = zero(progn_copy)
 
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> leapfrog_step(prog_new, progn_copy, x, dt, lf1, model), dprogn_copy, tend_copy)
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> leapfrog_step(prog_new, progn_copy, x, dt, lf1, model), dprogn_copy, tend_copy)
 
-    @test all(isapprox.(to_vec(fd_jvp[1])[1], to_vec(dtend)[1],rtol=1e-2,atol=1e-2))
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dtend)[1],rtol=1e-2,atol=1e-2))
 
     # 
     # single variable leapfrog step 
@@ -146,7 +176,7 @@
     w1 = L.robert_filter*L.williams_filter/2   
     w2 = L.robert_filter*(1-L.williams_filter)/2   
     @test all(dtendency .≈ dt*(1+w1-w2))
-    # d tend needs to be: dt* ( 1 + w1 - w2) (for every coefficient)
+    # ∂(tend) needs to be: dt* ( 1 + w1 - w2) (for every coefficient)
 
     function leapfrog(A_old, A_new, tendency, dt, lf, L)
         A_old_new = copy(A_old)
@@ -155,10 +185,10 @@
         return A_new_new
     end 
 
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x ->leapfrog(A_old_copy, A_new_copy, x, dt, lf1, L), one(tendency_copy), tendency_copy)
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x ->leapfrog(A_old_copy, A_new_copy, x, dt, lf1, L), one(tendency_copy), tendency_copy)
     
     # in this case it need to be dt*(1 - w2) (no contributation from A_old in FD)
-    @test all(isapprox.(fd_jvp[1], dt*(1-w2), rtol=1e-2))
+    @test all(isapprox.(fd_vjp[1], dt*(1-w2), rtol=1e-2))
 
     #
     # transform!(diagn, progn, lf2, model)
@@ -181,9 +211,9 @@
         return diag_copy
     end 
 
-    fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_diagn(diag_copy, x, lf2, model), ddiag_copy, progn_copy)
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_diagn(diag_copy, x, lf2, model), ddiag_copy, progn_copy)
     
-    @test all(isapprox.(to_vec(fd_jvp[1])[1], to_vec(dprogn)[1],rtol=1e-2,atol=1e-2))
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-2,atol=1e-2))
 
     # differnetiate wrt parameter 
     # write this as function (model, progn, diagn, 2\Delta t) -> progn_new
diff --git a/test/differentiability/speedy_transforms.jl b/test/differentiability/speedy_transforms.jl
index d8f761829..fa6e6bc75 100644
--- a/test/differentiability/speedy_transforms.jl
+++ b/test/differentiability/speedy_transforms.jl
@@ -27,8 +27,8 @@
             fill!(dspecs2, 1+1im)
 
             # finite difference comparision, seeded with a one adjoint to get the direct gradient
-            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dspecs2, grid)
-            @test isapprox(dgrid, fd_jvp[1])
+            fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dspecs2, grid)
+            @test isapprox(dgrid, fd_vjp[1])
 
             ## now backwards, as the input for spec we use the output of the forward transform
 
@@ -43,9 +43,9 @@
             dgrid2 = similar(grid)
             fill!(dgrid2, 1)
 
-            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dgrid2, specs)
+            fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dgrid2, specs)
 
-            @test isapprox(dspecs, fd_jvp[1])
+            @test isapprox(dspecs, fd_vjp[1])
         end 
 
         # test that d S^{-1}(S(x)) / dx = dx/dx = 1 (starting in both domains) 
@@ -95,8 +95,8 @@
         # The FD comparision passes, but it takes a long time to compute, so it's commented out. 
         #dgrid2 = similar(grid)
         #fill!(dgrid2, 1)
-        #fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_identity(x, S), dgrid2, grid)
-        #@test isapprox(dgrid, fd_jvp[1], rtol=0.01)
+        #fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_identity(x, S), dgrid2, grid)
+        #@test isapprox(dgrid, fd_vjp[1], rtol=0.01)
 
         # now start with spectral space, exclude for other grid because of https://github.com/SpeedyWeather/SpeedyWeather.jl/issues/626
         if fd_tests[i_grid]
@@ -171,9 +171,9 @@ end
             fill!(dcu2, 1)
 
             # finite difference comparision, seeded with a one adjoint to get the direct gradient
-            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> curl(x[1],x[2], S), dcu2, (u, v))
-            @test isapprox(du, fd_jvp[1][1], rtol=1e-6)
-            @test isapprox(dv, fd_jvp[1][2], rtol=1e-6)
+            fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> curl(x[1],x[2], S), dcu2, (u, v))
+            @test isapprox(du, fd_vjp[1][1], rtol=1e-6)
+            @test isapprox(dv, fd_vjp[1][2], rtol=1e-6)
 
             # div test
             u = transform(u_grid, S)
@@ -199,9 +199,9 @@ end
             ddiv2 = zero(ddiv)
             fill!(ddiv2, 1 + 1im)
 
-            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> divergence(x[1],x[2], S), ddiv2, (u, v))
-            @test isapprox(du, fd_jvp[1][1])
-            @test isapprox(dv, fd_jvp[1][2])
+            fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> divergence(x[1],x[2], S), ddiv2, (u, v))
+            @test isapprox(du, fd_vjp[1][1])
+            @test isapprox(dv, fd_vjp[1][2])
             @test sum(du) != 0 # nonzero gradient
             @test sum(dv) != 0 # nonzero gradient
 
@@ -230,8 +230,8 @@ end
             duv_input = zero(uv_input)
             duv_input = fill!(duv_input, 1+im)
 
-            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> uvfvor(x, S), duv_input, vor)
-            @test isapprox(dvor, fd_jvp[1])
+            fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> uvfvor(x, S), duv_input, vor)
+            @test isapprox(dvor, fd_vjp[1])
             @test sum(dvor) != 0 # nonzero gradient
 
             # UV_from_vordiv! 
@@ -261,9 +261,9 @@ end
             uv_input = zero(uv_input)
             duv_input = fill!(duv_input, 1+im)
 
-            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x-> uvfromvordiv(x[1], x[2], S), duv_input, (vor, div))
-            @test isapprox(dvor, fd_jvp[1][1][:,1]) 
-            @test isapprox(ddiv, fd_jvp[1][2][:,1])
+            fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x-> uvfromvordiv(x[1], x[2], S), duv_input, (vor, div))
+            @test isapprox(dvor, fd_vjp[1][1][:,1]) 
+            @test isapprox(ddiv, fd_vjp[1][2][:,1])
             @test sum(dvor) != 0 # nonzero gradient
             @test sum(ddiv) != 0 # nonzero gradient
 
@@ -278,9 +278,9 @@ end
             dres_∇2 = zero(res_∇)
             fill!(dres_∇2, 1+im)
 
-            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x-> ∇²(x, S), dres_∇2, vor)
+            fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x-> ∇²(x, S), dres_∇2, vor)
             @test sum(dvor) != 0 # non-zero 
-            @test isapprox(dvor, fd_jvp[1]) # and identical with FD
+            @test isapprox(dvor, fd_vjp[1]) # and identical with FD
 
             # test with the eigenvalues saved in S, result should just be seed * eigenvalues
             for i=1:(vor.m-1)
@@ -305,9 +305,9 @@ end
             dzonal_gradient2 = zero(dzonal_gradient)
             fill!(dzonal_gradient2, 1+im)
 
-            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x-> ∇(x, S), (dmerid_gradient2, dzonal_gradient2), vor)
+            fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x-> ∇(x, S), (dmerid_gradient2, dzonal_gradient2), vor)
             @test sum(dvor) != # nonzero 
-            @test isapprox(dvor, fd_jvp[1]) # and identical with FD
+            @test isapprox(dvor, fd_vjp[1]) # and identical with FD
         end 
     end 
 end 

From 922f0c16710da8392ba0e04986cb33575fdf69a5 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Tue, 28 Jan 2025 17:51:38 +0100
Subject: [PATCH 30/62] cleanup barotropic diff tests

---
 test/differentiability/barotropic.jl | 47 +++++++++-------------------
 1 file changed, 15 insertions(+), 32 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index 1b13b3f9d..e3b99ea91 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -5,6 +5,7 @@
     model = BarotropicModel(; spectral_grid)   # construct model
     simulation = initialize!(model)  
     initialize!(simulation)
+    run!(simulation, period=Day(5)) # spin-up to get nonzero values for all fields
 
     (; prognostic_variables, diagnostic_variables, model) = simulation
     (; Δt, Δt_millisec) = model.time_stepping
@@ -42,7 +43,8 @@
     # FD comparison 
     dprogn_2 = one(progn) # seed 
 
-    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
+    # for the full timestep, we need a bit higher precision 
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(15,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
     @test isapprox(to_vec(fd_vjp[1])[1], to_vec(d_progn)[1])
     
@@ -73,9 +75,9 @@
         return diagn_new
     end 
 
-    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> dynamics_tendencies(diagn_copy, x, lf2, model), ddiag_copy, progn_copy)
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(9,1), x -> dynamics_tendencies(diagn_copy, x, lf2, model), ddiag_copy, progn_copy)
 
-    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-2))
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
 
     # in the default configuration without forcing or drag, the barotropic model's don't dependent on the previous prognostic state 
     @test sum(to_vec(dprogn)[1]) ≈ 0 
@@ -94,14 +96,15 @@
 
     autodiff(Reverse, SpeedyWeather.horizontal_diffusion!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(model.horizontal_diffusion), Duplicated(model, make_zero(model)), Const(lf1))
 
-    function horizontal_diffusion(diagn, progn, diffusion, model, lf)
-        diagn_new = deepcopy(diagn)
-        SpeedyWeather.horizontal_diffusion!(diagn_new, progn, diffusion, model, lf)
-        return diagn_new
-    end 
+    # FD comparision not necessary, we have the exact values 
+    #function horizontal_diffusion(diagn, progn, diffusion, model, lf)
+    #    diagn_new = deepcopy(diagn)
+    #    SpeedyWeather.horizontal_diffusion!(diagn_new, progn, diffusion, model, lf)
+    #    return diagn_new
+    #end 
 
-    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> horizontal_diffusion(diagn_copy, x, model.horizontal_diffusion, model, lf1), ddiag_copy, progn_copy)
-    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-2))
+    #fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> horizontal_diffusion(diagn_copy, x, model.horizontal_diffusion, model, lf1), ddiag_copy, progn_copy)
+    #@test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-2))
 
     # ∂(progn)
     # should be row-wise `model.horizontal_diffusion.impl .* model.horizontal_diffusion.expl`
@@ -125,11 +128,6 @@
     lf1 = 2 
     lf2 = 2 
 
-    # set the tendencies back to zero for accumulation
-
-    # TENDENCIES, DIFFUSION, LEAPFROGGING AND TRANSFORM SPECTRAL STATE TO GRID
-    SpeedyWeather.horizontal_diffusion!(diagn, progn, model.horizontal_diffusion, model)
-
     progn_copy = deepcopy(progn)
     dprogn = one(progn_copy)
     dprogn_copy = one(progn_copy)
@@ -151,7 +149,7 @@
 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> leapfrog_step(prog_new, progn_copy, x, dt, lf1, model), dprogn_copy, tend_copy)
 
-    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dtend)[1],rtol=1e-2,atol=1e-2))
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dtend)[1],rtol=1e-5,atol=1e-5))
 
     # 
     # single variable leapfrog step 
@@ -178,18 +176,6 @@
     @test all(dtendency .≈ dt*(1+w1-w2))
     # ∂(tend) needs to be: dt* ( 1 + w1 - w2) (for every coefficient)
 
-    function leapfrog(A_old, A_new, tendency, dt, lf, L)
-        A_old_new = copy(A_old)
-        A_new_new = copy(A_new)
-        SpeedyWeather.leapfrog!(A_old_new, A_new_new, tendency, dt, lf, L)
-        return A_new_new
-    end 
-
-    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x ->leapfrog(A_old_copy, A_new_copy, x, dt, lf1, L), one(tendency_copy), tendency_copy)
-    
-    # in this case it need to be dt*(1 - w2) (no contributation from A_old in FD)
-    @test all(isapprox.(fd_vjp[1], dt*(1-w2), rtol=1e-2))
-
     #
     # transform!(diagn, progn, lf2, model)
     #
@@ -213,10 +199,7 @@
 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_diagn(diag_copy, x, lf2, model), ddiag_copy, progn_copy)
     
-    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-2,atol=1e-2))
-
-    # differnetiate wrt parameter 
-    # write this as function (model, progn, diagn, 2\Delta t) -> progn_new
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-3,atol=1e-3))
 end
 
 

From cd3e8378389342f9d06c8d1b46e66fd3df8e0b18 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Tue, 28 Jan 2025 18:04:05 +0100
Subject: [PATCH 31/62] more work on the correctness tests

---
 test/differentiability/barotropic.jl   |  5 ++---
 test/differentiability/timestepping.jl | 11 +++++------
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index e3b99ea91..389101d88 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -19,7 +19,6 @@
 
     d_progn = zero(progn)
     d_diag = make_zero(diagn)
-    d_model = make_zero(model)
 
     progn_new = zero(progn)
     dprogn_new = one(progn) # seed 
@@ -35,7 +34,7 @@
     dprogn_new = one(progn) # seed 
 
     # test that we can differentiate wrt an IC 
-    autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
+    autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, make_zero(model)))
 
     # nonzero gradient
     @test sum(to_vec(d_progn)[1]) != 0
@@ -46,7 +45,7 @@
     # for the full timestep, we need a bit higher precision 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(15,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
-    @test isapprox(to_vec(fd_vjp[1])[1], to_vec(d_progn)[1])
+    @test isapprox(to_vec(fd_vjp[1])[1], to_vec(d_progn)[1], rtol=1e-1) # we have to go really quite low with the tolerances here
     
     #
     # We go individually through all components of the time stepping and check 
diff --git a/test/differentiability/timestepping.jl b/test/differentiability/timestepping.jl
index 4da7cc9bd..66b5a2651 100644
--- a/test/differentiability/timestepping.jl
+++ b/test/differentiability/timestepping.jl
@@ -12,6 +12,7 @@
         model = model_type(; spectral_grid)   # construct model
         simulation = initialize!(model)  
         initialize!(simulation)
+        run!(simulation, period=Day(3))
 
         (; prognostic_variables, diagnostic_variables, model) = simulation
         (; Δt, Δt_millisec) = model.time_stepping
@@ -25,13 +26,12 @@
 
         d_progn = zero(progn)
         d_diag = make_zero(diagn)
-        d_model = make_zero(model)
 
         progn_new = zero(progn)
         dprogn_new = one(progn) # seed 
 
         # test that we can differentiate wrt an IC 
-        autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
+        autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, make_zero(model)))
 
         # nonzero gradient
         @test sum(to_vec(d_progn)[1]) != 0
@@ -39,9 +39,8 @@
         # FD comparison 
         dprogn_2 = one(progn) # seed 
 
-        fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
+        fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
-        @test isapprox(to_vec(fd_jvp[1])[1], to_vec(d_progn)[1])
+        @test isapprox(to_vec(fd_vjp[1])[1], to_vec(d_progn)[1])
     end 
-
-end
\ No newline at end of file
+end 
\ No newline at end of file

From 58899811a97f51883958dbdeeecc45007c98d661 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 29 Jan 2025 09:56:10 +0100
Subject: [PATCH 32/62] less hardcoded to_vec

---
 ext/SpeedyWeatherFiniteDifferencesExt.jl | 66 +++++++++++-------------
 1 file changed, 29 insertions(+), 37 deletions(-)

diff --git a/ext/SpeedyWeatherFiniteDifferencesExt.jl b/ext/SpeedyWeatherFiniteDifferencesExt.jl
index 6a5b5e6aa..477a54c93 100644
--- a/ext/SpeedyWeatherFiniteDifferencesExt.jl
+++ b/ext/SpeedyWeatherFiniteDifferencesExt.jl
@@ -49,12 +49,14 @@ end
 # that we don't want to be varied for our big data structures 
 function FiniteDifferences.to_vec(x::T) where {T <: Union{PrognosticVariables, PrognosticVariablesOcean, PrognosticVariablesLand, DiagnosticVariables, Tendencies, GridVariables, DynamicsVariables, PhysicsVariables, ParticleVariables}}
 
-    val_vecs_and_backs = map(name -> to_vec(getfield(x, name)), included_fields(T))
+    excluded_fields_pre, included_fields, excluded_fields_post = determine_included_fields(T)
+
+    val_vecs_and_backs = map(name -> to_vec(getfield(x, name)), included_fields)
     vals = first.(val_vecs_and_backs)
     backs = last.(val_vecs_and_backs)
 
-    vals_excluded_pre = map(name -> getfield(x, name), excluded_fields_pre(T))
-    vals_excluded_post = map(name -> getfield(x, name), excluded_fields_post(T))
+    vals_excluded_pre = map(name -> getfield(x, name), excluded_fields_pre)
+    vals_excluded_post = map(name -> getfield(x, name), excluded_fields_post)
 
     v, vals_from_vec = to_vec(vals)
     function structtype_from_vec(v::Vector{<:Real})
@@ -66,40 +68,30 @@ function FiniteDifferences.to_vec(x::T) where {T <: Union{PrognosticVariables, P
     return v, structtype_from_vec
 end
 
-included_fields(::Type{<:PrognosticVariables}) = (:vor, :div, :temp, :humid, :pres, :random_pattern, :ocean, :land, :tracers, :particles)
-excluded_fields_pre(::Type{<:PrognosticVariables}) = (:trunc, :nlat_half, :nlayers, :nparticles)
-excluded_fields_post(::Type{<:PrognosticVariables}) = (:scale, :clock)
-
-included_fields(::Type{<:PrognosticVariablesOcean}) = (:sea_surface_temperature, :sea_ice_concentration)
-excluded_fields_pre(::Type{<:PrognosticVariablesOcean}) = (:nlat_half, )
-excluded_fields_post(::Type{<:PrognosticVariablesOcean}) = ()
-
-included_fields(::Type{<:PrognosticVariablesLand}) = (:land_surface_temperature, :snow_depth, :soil_moisture_layer1, :soil_moisture_layer2)
-excluded_fields_pre(::Type{<:PrognosticVariablesLand}) = (:nlat_half, )
-excluded_fields_post(::Type{<:PrognosticVariablesLand}) = ()
-
-included_fields(::Type{<:DiagnosticVariables}) = (:tendencies, :grid, :dynamics, :physics, :particles, :column, :temp_average)
-excluded_fields_pre(::Type{<:DiagnosticVariables}) = (:trunc, :nlat_half, :nlayers, :nparticles)
-excluded_fields_post(::Type{<:DiagnosticVariables}) = (:scale,)
-
-included_fields(::Type{<:Tendencies}) = (:vor_tend, :div_tend, :temp_tend, :humid_tend, :u_tend, :v_tend, :pres_tend, :tracers_tend, :u_tend_grid, :v_tend_grid, :temp_tend_grid, :humid_tend_grid, :pres_tend_grid, :tracers_tend_grid)
-excluded_fields_pre(::Type{<:Tendencies}) = (:trunc, :nlat_half, :nlayers)
-excluded_fields_post(::Type{<:Tendencies}) = ()
-
-included_fields(::Type{<:GridVariables}) = (:vor_grid, :div_grid, :temp_grid, :temp_virt_grid, :humid_grid, :u_grid, :v_grid, :pres_grid, :tracers_grid, :random_pattern, :temp_grid_prev, :humid_grid_prev, :u_grid_prev, :v_grid_prev, :pres_grid_prev, :tracers_grid_prev)
-excluded_fields_pre(::Type{<:GridVariables}) = (:nlat_half, :nlayers)
-excluded_fields_post(::Type{<:GridVariables}) = ()
-
-included_fields(::Type{<:DynamicsVariables}) = (:a, :b, :a_grid, :b_grid, :a_2D, :b_2D, :a_2D_grid, :b_2D_grid, :uv∇lnp, :uv∇lnp_sum_above, :div_sum_above, :temp_virt, :geopot, :σ_tend, :∇lnp_x, :∇lnp_y, :u_mean_grid, :v_mean_grid, :div_mean_grid, :div_mean)
-excluded_fields_pre(::Type{<:DynamicsVariables}) = (:trunc, :nlat_half, :nlayers)
-excluded_fields_post(::Type{<:DynamicsVariables}) = ()
-
-included_fields(::Type{<:PhysicsVariables}) = (:precip_large_scale, :precip_convection, :precip_rate_large_scale, :precip_rate_convection, :cloud_top, :soil_moisture_availability, :sensible_heat_flux, :evaporative_flux, :surface_shortwave_up, :surface_shortwave_down, :surface_longwave_up, :surface_longwave_down, :outgoing_shortwave_radiation, :outgoing_longwave_radiation, :cos_zenith)
-excluded_fields_pre(::Type{<:PhysicsVariables}) = (:nlat_half,)
-excluded_fields_post(::Type{<:PhysicsVariables}) = ()
+function determine_included_fields(T::Type)
+    names = fieldnames(T)
+
+    included_field_types = Union{SpeedyWeather.AbstractDiagnosticVariables, 
+    SpeedyWeather.AbstractPrognosticVariables, SpeedyWeather.ColumnVariables,
+    NTuple, Dict{Symbol, <:Tuple}, Dict{Symbol, <:AbstractArray}, AbstractArray}
+
+    excluded_fields_pre = []
+    included_fields = []
+    excluded_fields_post = []
+
+    for name in names 
+        if fieldtype(T, name) <: included_field_types
+            push!(included_fields, name)
+        else 
+            if isempty(included_fields)
+                push!(excluded_fields_pre, name)
+            else 
+                push!(excluded_fields_post, name)
+            end 
+        end 
+    end 
 
-included_fields(::Type{<:ParticleVariables}) = (:locations, :u, :v, :σ_tend)
-excluded_fields_pre(::Type{<:ParticleVariables}) = (:nparticles, :nlat_half)
-excluded_fields_post(::Type{<:ParticleVariables}) = (:interpolator, )
+    return excluded_fields_pre, included_fields, excluded_fields_post
+end 
 
 end 
\ No newline at end of file

From bf1f842ff2addf6c2515d6a3111e05926c6dd50c Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 29 Jan 2025 09:57:14 +0100
Subject: [PATCH 33/62] remove flatten function

---
 test/differentiability/barotropic.jl |  6 +----
 test/differentiability/test_utils.jl | 39 +---------------------------
 2 files changed, 2 insertions(+), 43 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index 389101d88..802a49bca 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -25,11 +25,7 @@
 
     # test if differentiation works wrt copy! (there were some problems with it before)
     autodiff(Reverse, copy!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn))
-
-    # it's identity: 
-    @test all(flatten(d_progn)[1] .== 1)
-    @test all(flatten(d_progn)[2] .== 1)
-
+    
     progn_new = zero(progn)
     dprogn_new = one(progn) # seed 
 
diff --git a/test/differentiability/test_utils.jl b/test/differentiability/test_utils.jl
index fe44cbba1..50d252a2f 100644
--- a/test/differentiability/test_utils.jl
+++ b/test/differentiability/test_utils.jl
@@ -1,38 +1 @@
-# this is mainly there for testing purposes, and a bit more manual then above 
-function flatten(prog::PrognosticVariables{NF, ArrayType, NSTEPS, SpectralVariable2D, SpectralVariable3D, GridVariable2D}) where {NF,ArrayType, NSTEPS, SpectralVariable2D, SpectralVariable3D, GridVariable2D}
-
-    (; trunc, nlayers, nlat_half) = prog
-    nvars = 5 + length(progn.tracers)
-
-    # do it as a LTA 
-    prog_array = zeros(SpectralVariable3D, trunc+2, trunc+1, nlayers, NSTEPS, nvars)
-
-    for istep in 1:NSTEPS
-        prog_array[:, :, istep, 1] = prog.vor[istep]
-        prog_array[:, :, istep, 2] = prog.div[istep]
-        prog_array[:, :, istep, 3] = prog.temp[istep]
-        prog_array[:, :, istep, 4] = prog.humid[istep]
-        prog_array[:, 1, istep, 5] = prog.pres[istep]
-
-        for (i_key, (key, values)) in enumerate(prog.tracers)
-            prog_array[:,:, istep, 5+i_key] = values
-        end 
-    end 
-
-    nvars_grid = 6 
-    land_ocean_array = zeros(GridVariable2D, nlat_half, nvars_grid)
-    land_ocean_array[:, 1] = prog.ocean.sea_surface_temperature
-    land_ocean_array[:, 2] = prog.ocean.sea_ice_concentration
-    land_ocean_array[:, 3] = prog.land.land_surface_temperature
-    land_ocean_array[:, 4] = prog.land.snow_depth
-    land_ocean_array[:, 5] = prog.land.soil_moisture_layer1
-    land_ocean_array[:, 6] = prog.land.soil_moisture_layer2
-
-    return prog_array, land_ocean_array
-end 
-
-function Base.one(diag::DiagnosticVariables{NF}) where NF
-    vec, re = to_vec(diag)
-    vec .= NF(1)
-    return re(vec)
-end 
\ No newline at end of file
+# additional utitility functions 
\ No newline at end of file

From 26141d98a188ecd762090148d0055223c3c24b21 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 29 Jan 2025 10:34:20 +0100
Subject: [PATCH 34/62] adjust new surface fluxes, less type instabilities

---
 src/physics/surface_fluxes/heat.jl     | 12 ++++++------
 src/physics/surface_fluxes/moisture.jl | 12 ++++++------
 src/physics/surface_fluxes/momentum.jl |  6 +++---
 3 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/src/physics/surface_fluxes/heat.jl b/src/physics/surface_fluxes/heat.jl
index 2fb38e989..2cceb3e88 100644
--- a/src/physics/surface_fluxes/heat.jl
+++ b/src/physics/surface_fluxes/heat.jl
@@ -49,10 +49,10 @@ SurfaceOceanHeatFlux(SG::SpectralGrid; kwargs...) = SurfaceOceanHeatFlux{SG.NF}(
 initialize!(::SurfaceOceanHeatFlux, ::PrimitiveEquation) = nothing
 
 function surface_heat_flux!(
-    column::ColumnVariables,
+    column::ColumnVariables{NF},
     heat_flux::SurfaceOceanHeatFlux,
     model::PrimitiveEquation,
-)
+) where NF
     cₚ = model.atmosphere.heat_capacity
     (; heat_exchange) = heat_flux
 
@@ -67,7 +67,7 @@ function surface_heat_flux!(
 
     # SPEEDY documentation Eq. 54/56, land/sea fraction included
     # Only flux from sea if available (not NaN) otherwise zero flux
-    flux_ocean  = isfinite(T_skin_ocean) ? ρ*drag_ocean*V₀*cₚ*(T_skin_ocean  - T)*(1-land_fraction) : 0
+    flux_ocean  = isfinite(T_skin_ocean) ? ρ*drag_ocean*V₀*cₚ*(T_skin_ocean  - T)*(1-land_fraction) : zero(NF)
     column.flux_temp_upward[end] += flux_ocean
     column.sensible_heat_flux = flux_ocean      # ocean sets the flux (=), land accumulates (+=)
 
@@ -89,10 +89,10 @@ SurfaceLandHeatFlux(SG::SpectralGrid; kwargs...) = SurfaceLandHeatFlux{SG.NF}(;
 initialize!(::SurfaceLandHeatFlux, ::PrimitiveEquation) = nothing
 
 function surface_heat_flux!(
-    column::ColumnVariables,
+    column::ColumnVariables{NF},
     heat_flux::SurfaceLandHeatFlux,
     model::PrimitiveEquation,
-)
+) where NF
     cₚ = model.atmosphere.heat_capacity
     (; heat_exchange) = heat_flux
 
@@ -107,7 +107,7 @@ function surface_heat_flux!(
 
     # SPEEDY documentation Eq. 54/56, land/sea fraction included
     # Only flux from sea if available (not NaN) otherwise zero flux
-    flux_land  = isfinite(T_skin_land) ? ρ*drag_land*V₀*cₚ*(T_skin_land  - T)*land_fraction : 0
+    flux_land  = isfinite(T_skin_land) ? ρ*drag_land*V₀*cₚ*(T_skin_land  - T)*land_fraction : zero(NF)
     column.flux_temp_upward[end] += flux_land
     column.sensible_heat_flux += flux_land      # ocean sets the flux (=), land accumulates (+=)
 
diff --git a/src/physics/surface_fluxes/moisture.jl b/src/physics/surface_fluxes/moisture.jl
index b3a8795fe..1ac151992 100644
--- a/src/physics/surface_fluxes/moisture.jl
+++ b/src/physics/surface_fluxes/moisture.jl
@@ -49,10 +49,10 @@ SurfaceOceanEvaporation(SG::SpectralGrid; kwargs...) = SurfaceOceanEvaporation{S
 initialize!(::SurfaceOceanEvaporation, ::PrimitiveWet) = nothing
 
 function surface_evaporation!(
-    column::ColumnVariables,
+    column::ColumnVariables{NF},
     evaporation::SurfaceOceanEvaporation,
     model::PrimitiveWet,
-)
+) where NF
 
     (; skin_temperature_sea, pres) = column
     (; moisture_exchange) = evaporation
@@ -70,7 +70,7 @@ function surface_evaporation!(
     drag_sea = evaporation.use_boundary_layer_drag ? column.boundary_layer_drag : moisture_exchange
 
     # SPEEDY documentation eq. 55/57
-    flux_sea = isfinite(skin_temperature_sea) ? ρ*drag_sea*V₀*max(sat_humid_sea  - surface_humid, 0)*(1-land_fraction) : 0
+    flux_sea = isfinite(skin_temperature_sea) ? ρ*drag_sea*V₀*max(sat_humid_sea  - surface_humid, zero(NF))*(1-land_fraction) : zero(NF)
     column.flux_humid_upward[end] += flux_sea
     column.evaporative_flux = flux_sea
 
@@ -92,10 +92,10 @@ SurfaceLandEvaporation(SG::SpectralGrid; kwargs...) = SurfaceLandEvaporation{SG.
 initialize!(::SurfaceLandEvaporation, ::PrimitiveWet) = nothing
 
 function surface_evaporation!(
-    column::ColumnVariables,
+    column::ColumnVariables{NF},
     evaporation::SurfaceLandEvaporation,
     model::PrimitiveWet,
-)
+) where NF
 
     (; skin_temperature_land, pres) = column
     (; moisture_exchange) = evaporation
@@ -114,7 +114,7 @@ function surface_evaporation!(
     drag_land = evaporation.use_boundary_layer_drag ? column.boundary_layer_drag : moisture_exchange
 
     # SPEEDY documentation eq. 55/57
-    flux_land = isfinite(skin_temperature_land) && isfinite(α) ? ρ*drag_land*V₀*max(α*sat_humid_land  - surface_humid, 0)*land_fraction : 0
+    flux_land = isfinite(skin_temperature_land) && isfinite(α) ? ρ*drag_land*V₀*max(α*sat_humid_land  - surface_humid, zero(NF))*land_fraction : zero(NF)
     column.flux_humid_upward[end] += flux_land      # end=lowermost layer
     column.evaporative_flux += flux_land            # ocean sets the flux (=), land accumulates (+=)
 
diff --git a/src/physics/surface_fluxes/momentum.jl b/src/physics/surface_fluxes/momentum.jl
index 0c1392c9d..93b871797 100644
--- a/src/physics/surface_fluxes/momentum.jl
+++ b/src/physics/surface_fluxes/momentum.jl
@@ -25,9 +25,9 @@ end
 SurfaceWind(SG::SpectralGrid; kwargs...) = SurfaceWind{SG.NF}(; kwargs...)
 initialize!(::SurfaceWind, ::PrimitiveEquation) = nothing
 
-function surface_wind_stress!(  column::ColumnVariables,
+function surface_wind_stress!(  column::ColumnVariables{NF},
                                 surface_wind::SurfaceWind,
-                                model::PrimitiveEquation)
+                                model::PrimitiveEquation) where NF
 
     (; land_fraction) = column
     (; f_wind, V_gust, drag_land, drag_sea) = surface_wind
@@ -48,7 +48,7 @@ function surface_wind_stress!(  column::ColumnVariables,
     # surface wind stress: quadratic drag, fractional land-sea mask
     ρ = column.surface_air_density
     V₀ = column.surface_wind_speed
-    drag = land_fraction*drag_land + (1-land_fraction)*drag_sea
+    drag = land_fraction*drag_land + (one(NF)-land_fraction)*drag_sea
 
     # SPEEDY documentation eq. 52, 53, accumulate fluxes with +=
     column.flux_u_upward[end] -= ρ*drag*V₀*surface_u

From 50dd2eca52fa0685cadbb64aad941aadc291fd0d Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 29 Jan 2025 10:44:21 +0100
Subject: [PATCH 35/62] adjust show and copy! progn

---
 src/dynamics/prognostic_variables.jl | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl
index 9c570f259..03d072e88 100644
--- a/src/dynamics/prognostic_variables.jl
+++ b/src/dynamics/prognostic_variables.jl
@@ -169,12 +169,16 @@ function Base.show(
     println(io, "├ random_pattern: T$trunc, 1-layer LowerTriangularArray{$NF}")
     println(io, "├┐ocean: PrognosticVariablesOcean{$NF}")
     println(io, "│├ sea_surface_temperature:  $nlat-ring $Grid")
-    println(io, "│└ sea_ice_concentration:    $nlat-ring $Grid")
+    println(io, "│├ sea_ice_concentration:    $nlat-ring $Grid")
+    println(io, "│├ sensible_heat_flux:       $nlat-ring $Grid")
+    println(io, "│└ evaporative_flux:         $nlat-ring $Grid")
     println(io, "├┐land:  PrognosticVariablesLand{$NF}")
     println(io, "│├ land_surface_temperature: $nlat-ring $Grid")
     println(io, "│├ snow_depth: $nlat-ring $Grid")
     println(io, "│├ soil_moisture_layer1:     $nlat-ring $Grid")
-    println(io, "│└ soil_moisture_layer2:     $nlat-ring $Grid")
+    println(io, "│├ soil_moisture_layer2:     $nlat-ring $Grid")
+    println(io, "│├ sensible_heat_flux:       $nlat-ring $Grid")
+    println(io, "│└ evaporative_flux:         $nlat-ring $Grid")
     println(io, "├ tracers: $(length(tracer_names)), $tracer_names")
     println(io, "├ particles: $nparticles-element $(typeof(progn.particles))")
     println(io, "├ scale: $(progn.scale[])")
@@ -198,12 +202,16 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
     # ocean
     progn_new.ocean.sea_surface_temperature .= progn_old.ocean.sea_surface_temperature
     progn_new.ocean.sea_ice_concentration .= progn_old.ocean.sea_ice_concentration
-    
+    progn_new.ocean.sensible_heat_flux = progn_old.ocean.sensible_heat_flux
+    progn_new.ocean.evaporative_flux = progn_old.ocean.evaporative_flux
+
     # land
     progn_new.land.land_surface_temperature .= progn_old.land.land_surface_temperature
     progn_new.land.snow_depth .= progn_old.land.snow_depth
     progn_new.land.soil_moisture_layer1 .= progn_old.land.soil_moisture_layer1
     progn_new.land.soil_moisture_layer2 .= progn_old.land.soil_moisture_layer2
+    progn_new.land.sensible_heat_flux = progn_old.land.sensible_heat_flux
+    progn_new.land.evaporative_flux = progn_old.land.evaporative_flux
 
     # copy over tracers
     for (key, value) in progn_old.tracers
@@ -264,11 +272,16 @@ function Base.fill!(progn::PrognosticVariables{NF}, value::Number) where NF
     # ocean
     progn.ocean.sea_surface_temperature .= value_NF
     progn.ocean.sea_ice_concentration .= value_NF
+    progn.ocean.sensible_heat_flux .= value_NF
+    progn.ocean.evaporative_flux .= value_NF
+
     # land
     progn.land.land_surface_temperature .= value_NF
     progn.land.snow_depth .= value_NF
     progn.land.soil_moisture_layer1 .= value_NF
     progn.land.soil_moisture_layer2 .= value_NF
+    progn.land.sensible_heat_flux .= value_NF
+    progn.land.evaporative_flux .= value_NF
 
     # fill tracers
     for (key, value) in progn.tracers 

From a3f16382dd6b56cf79b0c38a40804152c86af26e Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 29 Jan 2025 10:46:45 +0100
Subject: [PATCH 36/62] fix copy progn

---
 src/dynamics/prognostic_variables.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl
index 03d072e88..4a43845e0 100644
--- a/src/dynamics/prognostic_variables.jl
+++ b/src/dynamics/prognostic_variables.jl
@@ -202,16 +202,16 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
     # ocean
     progn_new.ocean.sea_surface_temperature .= progn_old.ocean.sea_surface_temperature
     progn_new.ocean.sea_ice_concentration .= progn_old.ocean.sea_ice_concentration
-    progn_new.ocean.sensible_heat_flux = progn_old.ocean.sensible_heat_flux
-    progn_new.ocean.evaporative_flux = progn_old.ocean.evaporative_flux
+    progn_new.ocean.sensible_heat_flux .= progn_old.ocean.sensible_heat_flux
+    progn_new.ocean.evaporative_flux .= progn_old.ocean.evaporative_flux
 
     # land
     progn_new.land.land_surface_temperature .= progn_old.land.land_surface_temperature
     progn_new.land.snow_depth .= progn_old.land.snow_depth
     progn_new.land.soil_moisture_layer1 .= progn_old.land.soil_moisture_layer1
     progn_new.land.soil_moisture_layer2 .= progn_old.land.soil_moisture_layer2
-    progn_new.land.sensible_heat_flux = progn_old.land.sensible_heat_flux
-    progn_new.land.evaporative_flux = progn_old.land.evaporative_flux
+    progn_new.land.sensible_heat_flux .= progn_old.land.sensible_heat_flux
+    progn_new.land.evaporative_flux .= progn_old.land.evaporative_flux
 
     # copy over tracers
     for (key, value) in progn_old.tracers

From 4dc8944406199bbe53111838d624f8c08bedb516 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 29 Jan 2025 10:48:14 +0100
Subject: [PATCH 37/62] remove merge conflict message from changelog

---
 CHANGELOG.md | 1 -
 1 file changed, 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ed036f609..46c9986a4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,7 +9,6 @@
 
 - CompatHelper: Allow for Makie.jl v0.22 [#663](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/663)
 - Prescribed ocean/land heat fluxes modular [#667](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/667)
->>>>>>> origin/main
 - NetCDF output for tracers defined, precipitation rate output at initialize! again [#657](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/657)
 - Precipitation rate from large-scale condensation/convection for coupling [#669](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/669)
 - Fixed a bug, so that sum now works for LowerTraingularArrays [#668](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/668)

From ece44499594b535c71985fd5572b35a1d3f93bb1 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 29 Jan 2025 12:28:23 +0100
Subject: [PATCH 38/62] adjust CI test

---
 src/dynamics/prognostic_variables.jl | 12 ++++++------
 test/spectral_transform_ad_rules.jl  |  8 +-------
 2 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl
index 4a43845e0..3e2b8157a 100644
--- a/src/dynamics/prognostic_variables.jl
+++ b/src/dynamics/prognostic_variables.jl
@@ -174,7 +174,7 @@ function Base.show(
     println(io, "│└ evaporative_flux:         $nlat-ring $Grid")
     println(io, "├┐land:  PrognosticVariablesLand{$NF}")
     println(io, "│├ land_surface_temperature: $nlat-ring $Grid")
-    println(io, "│├ snow_depth: $nlat-ring $Grid")
+    println(io, "│├ snow_depth:               $nlat-ring $Grid")
     println(io, "│├ soil_moisture_layer1:     $nlat-ring $Grid")
     println(io, "│├ soil_moisture_layer2:     $nlat-ring $Grid")
     println(io, "│├ sensible_heat_flux:       $nlat-ring $Grid")
@@ -191,6 +191,11 @@ end
 Copies entries of `progn_old` into `progn_new`."""
 function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariables)
 
+    progn_new.random_pattern .= progn_old.random_pattern
+
+    copy!(progn_new.clock, progn_old.clock)
+    progn_new.scale[] = progn_old.scale[]
+
     for i in eachindex(progn_new.vor)   # each leapfrog time step
         progn_new.vor[i] .= progn_old.vor[i]
         progn_new.div[i] .= progn_old.div[i]
@@ -229,11 +234,6 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
         progn_new.particles .= progn_old.particles
     end
 
-    progn_new.random_pattern .= progn_old.random_pattern
-
-    copy!(progn_new.clock, progn_old.clock)
-    progn_new.scale[] = progn_old.scale[]
-
     return progn_new
 end
 
diff --git a/test/spectral_transform_ad_rules.jl b/test/spectral_transform_ad_rules.jl
index 9bf393c68..611fc903c 100644
--- a/test/spectral_transform_ad_rules.jl
+++ b/test/spectral_transform_ad_rules.jl
@@ -63,12 +63,6 @@ end
     # in a seperate test set. But we do want to ensure in the regular CI that 
     # we don't commit some kind of problem for the Enzyme differentiability
     # so, we test here if we get a non-zero gradient from the timestepping.  
-    function timestep_oop!(progn_new::PrognosticVariables, progn_old::PrognosticVariables, diagn, dt, model)
-        copy!(progn_new, progn_old)
-        SpeedyWeather.timestep!(progn_new, diagn, dt, model)
-        return nothing
-    end 
-
     spectral_grid = SpectralGrid(trunc=15, nlayers=3)          # define resolution
     model = PrimitiveWetModel(; spectral_grid)   # construct model
     simulation = initialize!(model)  
@@ -91,6 +85,6 @@ end
     progn_new = zero(progn)
     dprogn_new = one(progn) # seed 
 
-    autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
+    autodiff(Reverse, timestep!, Const, Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
     @test sum(to_vec(d_progn)[1]) != 0
 end 
\ No newline at end of file

From 65823f05b7b982d573ac457fe097acfe3f8a37ca Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 29 Jan 2025 12:39:37 +0100
Subject: [PATCH 39/62] fix CI test (again)

---
 test/spectral_transform_ad_rules.jl | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/test/spectral_transform_ad_rules.jl b/test/spectral_transform_ad_rules.jl
index 611fc903c..c9030daca 100644
--- a/test/spectral_transform_ad_rules.jl
+++ b/test/spectral_transform_ad_rules.jl
@@ -67,7 +67,8 @@ end
     model = PrimitiveWetModel(; spectral_grid)   # construct model
     simulation = initialize!(model)  
     initialize!(simulation)
-
+    run!(simulation, period=Day(1))
+    
     (; prognostic_variables, diagnostic_variables, model) = simulation
     (; Δt, Δt_millisec) = model.time_stepping
     dt = 2Δt
@@ -85,6 +86,6 @@ end
     progn_new = zero(progn)
     dprogn_new = one(progn) # seed 
 
-    autodiff(Reverse, timestep!, Const, Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
+    autodiff(Reverse, SpeedyWeather.timestep!, Const, Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
     @test sum(to_vec(d_progn)[1]) != 0
 end 
\ No newline at end of file

From 1375887abc23c26dcfe62baf7cb0e047465afb7f Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 29 Jan 2025 17:40:25 +0100
Subject: [PATCH 40/62] copy! progn avoids ringgrids broadcasting

---
 src/dynamics/prognostic_variables.jl | 33 ++++++++++++++--------------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl
index 3e2b8157a..90e15943a 100644
--- a/src/dynamics/prognostic_variables.jl
+++ b/src/dynamics/prognostic_variables.jl
@@ -191,11 +191,6 @@ end
 Copies entries of `progn_old` into `progn_new`."""
 function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariables)
 
-    progn_new.random_pattern .= progn_old.random_pattern
-
-    copy!(progn_new.clock, progn_old.clock)
-    progn_new.scale[] = progn_old.scale[]
-
     for i in eachindex(progn_new.vor)   # each leapfrog time step
         progn_new.vor[i] .= progn_old.vor[i]
         progn_new.div[i] .= progn_old.div[i]
@@ -204,19 +199,20 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
         progn_new.pres[i] .= progn_old.pres[i]
     end
 
+    # TO-DO: currently there are some problems with Enzyme and RingGrids broadcasting
     # ocean
-    progn_new.ocean.sea_surface_temperature .= progn_old.ocean.sea_surface_temperature
-    progn_new.ocean.sea_ice_concentration .= progn_old.ocean.sea_ice_concentration
-    progn_new.ocean.sensible_heat_flux .= progn_old.ocean.sensible_heat_flux
-    progn_new.ocean.evaporative_flux .= progn_old.ocean.evaporative_flux
+    progn_new.ocean.sea_surface_temperature.data .= progn_old.ocean.sea_surface_temperature.data
+    progn_new.ocean.sea_ice_concentration.data .= progn_old.ocean.sea_ice_concentration.data
+    progn_new.ocean.sensible_heat_flux.data .= progn_old.ocean.sensible_heat_flux.data
+    progn_new.ocean.evaporative_flux.data .= progn_old.ocean.evaporative_flux.data
 
     # land
-    progn_new.land.land_surface_temperature .= progn_old.land.land_surface_temperature
-    progn_new.land.snow_depth .= progn_old.land.snow_depth
-    progn_new.land.soil_moisture_layer1 .= progn_old.land.soil_moisture_layer1
-    progn_new.land.soil_moisture_layer2 .= progn_old.land.soil_moisture_layer2
-    progn_new.land.sensible_heat_flux .= progn_old.land.sensible_heat_flux
-    progn_new.land.evaporative_flux .= progn_old.land.evaporative_flux
+    progn_new.land.land_surface_temperature.data .= progn_old.land.land_surface_temperature.data
+    progn_new.land.snow_depth.data .= progn_old.land.snow_depth.data
+    progn_new.land.soil_moisture_layer1.data .= progn_old.land.soil_moisture_layer1.data
+    progn_new.land.soil_moisture_layer2.data .= progn_old.land.soil_moisture_layer2.data
+    progn_new.land.sensible_heat_flux.data .= progn_old.land.sensible_heat_flux.data
+    progn_new.land.evaporative_flux.data .= progn_old.land.evaporative_flux.data
 
     # copy over tracers
     for (key, value) in progn_old.tracers
@@ -234,7 +230,12 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
         progn_new.particles .= progn_old.particles
     end
 
-    return progn_new
+    progn_new.random_pattern.data .= progn_old.random_pattern.data
+
+    copy!(progn_new.clock, progn_old.clock)
+    progn_new.scale[] = progn_old.scale[]
+
+    return nothing
 end
 
 function Base.zero(progn::PrognosticVariables{NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}) where {NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}

From 67e0dac0202fec729f6904a45bc0f62c0f9a1e1e Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Thu, 30 Jan 2025 14:21:26 +0100
Subject: [PATCH 41/62] barotropic correctness tests done

---
 test/differentiability/barotropic.jl     | 11 ++++++++---
 test/differentiability/timestep_utils.jl |  8 ++++----
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index 802a49bca..d6f14bcad 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -6,7 +6,8 @@
     simulation = initialize!(model)  
     initialize!(simulation)
     run!(simulation, period=Day(5)) # spin-up to get nonzero values for all fields
-
+    initialize!(simulation; period=Day(1))
+    
     (; prognostic_variables, diagnostic_variables, model) = simulation
     (; Δt, Δt_millisec) = model.time_stepping
     dt = 2Δt
@@ -14,6 +15,9 @@
     progn = prognostic_variables
     diagn = diagnostic_variables
 
+    # TO-DO: The first time we execute this, the gradient is different. Why?
+    timestep_oop!(make_zero(progn), progn, diagn, dt, model)
+
     diagn_copy = deepcopy(diagn)
     progn_copy = deepcopy(progn)
 
@@ -25,7 +29,7 @@
 
     # test if differentiation works wrt copy! (there were some problems with it before)
     autodiff(Reverse, copy!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn))
-    
+
     progn_new = zero(progn)
     dprogn_new = one(progn) # seed 
 
@@ -41,7 +45,8 @@
     # for the full timestep, we need a bit higher precision 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(15,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
-    @test isapprox(to_vec(fd_vjp[1])[1], to_vec(d_progn)[1], rtol=1e-1) # we have to go really quite low with the tolerances here
+    @test isapprox(to_vec(fd_vjp[1])[1], to_vec(d_progn)[1], rtol=0.05) # we have to go really quite low with the tolerances here
+    @test mean(abs.(to_vec(fd_vjp[1])[1] - to_vec(d_progn)[1])) < 0.002
     
     #
     # We go individually through all components of the time stepping and check 
diff --git a/test/differentiability/timestep_utils.jl b/test/differentiability/timestep_utils.jl
index 0986058a7..090a51940 100644
--- a/test/differentiability/timestep_utils.jl
+++ b/test/differentiability/timestep_utils.jl
@@ -1,11 +1,11 @@
-function timestep_oop!(progn_new::PrognosticVariables, progn_old::PrognosticVariables, diagn, dt, model)
+function timestep_oop!(progn_new::PrognosticVariables, progn_old::PrognosticVariables, diagn, dt, model, lf1=2, lf2=2)
     copy!(progn_new, progn_old)
-    SpeedyWeather.timestep!(progn_new, diagn, dt, model)
+    SpeedyWeather.timestep!(progn_new, diagn, dt, model, lf1, lf2)
     return nothing
 end 
 
-function timestep_oop(progn, diagn, dt, model)
+function timestep_oop(progn, diagn, dt, model, lf1=2, lf2=2)
     progn_new = zero(progn)
-    timestep_oop!(progn_new, progn, diagn, dt, model)
+    timestep_oop!(progn_new, progn, diagn, dt, model, lf1, lf2)
     return progn_new
 end 

From b27609e30f8073716baca20cf0c31a250fe74cf5 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Thu, 30 Jan 2025 15:37:51 +0100
Subject: [PATCH 42/62] WIP tests

---
 test/differentiability/runtests.jl          | 4 ----
 test/differentiability/speedy_transforms.jl | 4 +++-
 test/differentiability/test_utils.jl        | 8 +++++++-
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/test/differentiability/runtests.jl b/test/differentiability/runtests.jl
index 992a41251..13881c875 100644
--- a/test/differentiability/runtests.jl
+++ b/test/differentiability/runtests.jl
@@ -6,10 +6,6 @@ import EnzymeTestUtils: test_approx
 import FiniteDifferences: j′vp, grad, central_fdm
 import AbstractFFTs
 
-grid_types = [FullGaussianGrid, OctahedralGaussianGrid] # one full and one reduced grid, both Gaussian to have exact transforms 
-grid_dealiasing = [2, 3]
-fd_tests = [true, true] 
-
 # UTILITIES 
 include("test_utils.jl")
 include("timestep_utils.jl")
diff --git a/test/differentiability/speedy_transforms.jl b/test/differentiability/speedy_transforms.jl
index fa6e6bc75..ab09552ec 100644
--- a/test/differentiability/speedy_transforms.jl
+++ b/test/differentiability/speedy_transforms.jl
@@ -1,5 +1,7 @@
 # Tests for SpeedyTransforms
-
+grid_types = [FullGaussianGrid, OctahedralGaussianGrid] # one full and one reduced grid, both Gaussian to have exact transforms 
+grid_dealiasing = [2, 3]
+fd_tests = [true, true] 
 @testset "Differentiability: Complete Transform Enzyme" begin
     # make a high level finite difference test of the whole transform
     # can't use Enzyme or ChainRule Test tools for tests for that
diff --git a/test/differentiability/test_utils.jl b/test/differentiability/test_utils.jl
index 50d252a2f..7476ef4f7 100644
--- a/test/differentiability/test_utils.jl
+++ b/test/differentiability/test_utils.jl
@@ -1 +1,7 @@
-# additional utitility functions 
\ No newline at end of file
+# additional utitility functions 
+
+function Base.one(diag::DiagnosticVariables{NF}) where NF
+    vec, re = to_vec(diag)
+    vec .= NF(1)
+    return re(vec)
+end 
\ No newline at end of file

From 525ef09ca4525b34b804906bab5ce109370892e1 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Thu, 30 Jan 2025 17:23:29 +0100
Subject: [PATCH 43/62] WIP primtive wet tests

---
 test/differentiability/barotropic.jl   |   6 +-
 test/differentiability/primitivewet.jl | 260 +++++++++++++++++++++++++
 2 files changed, 263 insertions(+), 3 deletions(-)
 create mode 100644 test/differentiability/primitivewet.jl

diff --git a/test/differentiability/barotropic.jl b/test/differentiability/barotropic.jl
index d6f14bcad..a46ecc081 100644
--- a/test/differentiability/barotropic.jl
+++ b/test/differentiability/barotropic.jl
@@ -1,5 +1,5 @@
 ### Experiments going a bit deeper into the timestepping of the barotropic model
-@testset "Differentiability: Barotropic Model Timestepping" begin 
+@testset "Differentiability: Barotropic Model Components" begin 
     # T15 still yields somewhat sensible dynamics, that's why it's chosen here
     spectral_grid = SpectralGrid(trunc=15, nlayers=1)          # define resolution
     model = BarotropicModel(; spectral_grid)   # construct model
@@ -46,8 +46,8 @@
     fd_vjp = FiniteDifferences.j′vp(central_fdm(15,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
     @test isapprox(to_vec(fd_vjp[1])[1], to_vec(d_progn)[1], rtol=0.05) # we have to go really quite low with the tolerances here
-    @test mean(abs.(to_vec(fd_vjp[1])[1] - to_vec(d_progn)[1])) < 0.002
-    
+    @test mean(abs.(to_vec(fd_vjp[1])[1] - to_vec(d_progn)[1])) < 0.002 # so we check a few extra statistics
+    @test maximum(to_vec(fd_vjp[1].vor)[1] - to_vec(d_progn.vor)[1]) < 0.05
     #
     # We go individually through all components of the time stepping and check 
     # correctness
diff --git a/test/differentiability/primitivewet.jl b/test/differentiability/primitivewet.jl
new file mode 100644
index 000000000..4721a7a2e
--- /dev/null
+++ b/test/differentiability/primitivewet.jl
@@ -0,0 +1,260 @@
+### Experiments going a bit deeper into the timestepping of the barotropic model
+@testset "Differentiability: Primitive Wet Model Components" begin 
+    # T15 still yields somewhat sensible dynamics, that's why it's chosen here
+    spectral_grid = SpectralGrid(trunc=15, nlayers=3)          # define resolution
+    model = PrimitiveWetModel(; spectral_grid)   # construct model
+    simulation = initialize!(model)  
+    initialize!(simulation)
+    run!(simulation, period=Day(5)) # spin-up to get nonzero values for all fields
+    initialize!(simulation; period=Day(1))
+    
+    (; prognostic_variables, diagnostic_variables, model) = simulation
+    (; Δt, Δt_millisec) = model.time_stepping
+    dt = 2Δt
+
+    progn = prognostic_variables
+    diagn = diagnostic_variables
+
+    # TO-DO: The first time we execute this, the gradient is different. Why?
+    timestep_oop!(make_zero(progn), progn, diagn, dt, model)
+
+    #
+    # We go individually through all components of the time stepping and check 
+    # correctness
+    #
+
+    fill!(diagn.tendencies, 0, PrimitiveWetModel)
+    (; time) = progn.clock 
+
+    #
+    # model physics 
+    # 
+    progn_copy = deepcopy(progn)
+    dprogn = one(progn)
+    ddiagn = one(diagn)
+    ddiagn_copy = deepcopy(ddiagn)
+
+    autodiff(Reverse, SpeedyWeather.parameterization_tendencies!, Const, Duplicated(diagn, ddiagn), Duplicated(progn, dprogn), Const(progn.clock.time), Duplicated(model, make_zero(model)))
+
+    function parameterization_tendencies(diagn, progn, time, model)
+        diagn_new = deepcopy(diagn)
+        SpeedyWeather.parameterization_tendencies!(diagn_new, progn, time, model)
+        return diagn_new
+    end 
+
+    # TODO: non-working at the moment -> NaN
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> parameterization_tendencies(diagn_copy, x, progn.clock.time, model), ddiagn_copy, progn_copy)
+    #@test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
+    
+    #
+    # ocean 
+    # 
+
+    progn_copy = deepcopy(progn)
+    dprogn = one(progn)
+    ddiagn = make_zero(diagn)
+    dprogn_copy = deepcopy(dprogn)
+    diagn_copy = deepcopy(diagn)
+
+    autodiff(Reverse, SpeedyWeather.ocean_timestep!, Const, Duplicated(progn, dprogn), Duplicated(diagn, ddiagn), Duplicated(model, make_zero(model)))
+
+    function ocean_timestep(progn, diagn, model)
+        progn_new = deepcopy(progn)
+        SpeedyWeather.ocean_timestep!(progn_new, diagn, model)
+        return progn_new
+    end 
+
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> ocean_timestep(progn_copy, x, model), dprogn_copy, diagn_copy)
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
+
+    #
+    # land 
+    # 
+    
+    progn_copy = deepcopy(progn)
+    dprogn = one(progn)
+    ddiagn = make_zero(diagn)
+    dprogn_copy = deepcopy(dprogn)
+    diagn_copy = deepcopy(diagn)
+
+    autodiff(Reverse, SpeedyWeather.land_timestep!, Const, Duplicated(progn, dprogn), Duplicated(diagn, ddiagn), Duplicated(model, make_zero(model)))
+
+    function land_timestep(progn, diagn, model)
+        progn_new = deepcopy(progn)
+        SpeedyWeather.ocean_timestep!(progn_new, diagn, model)
+        return progn_new
+    end 
+
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> land_timestep(progn_copy, x, model), dprogn_copy, diagn_copy)
+    
+    #####
+    # DYNAMICS 
+    lf2 = 2 
+    #
+    # drag! 
+    #
+
+    diagn_copy = deepcopy(diagn)
+    ddiag = one(diagn_copy)
+    ddiag_copy = deepcopy(ddiag)
+    progn_copy = deepcopy(progn)
+    dprogn = make_zero(progn)
+
+    autodiff(Reverse, SpeedyWeather.drag!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(lf2), Duplicated(model, make_zero(model)))
+
+    function drag(diagn, progn, lf, model)
+        diagn_new = deepcopy(diagn)
+        SpeedyWeather.drag!(diagn_new, progn, lf, model)
+        return diagn_new
+    end 
+
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(9,1), x -> drag(diagn_copy, x, lf2, model), ddiag_copy, progn_copy)
+
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
+
+    # in the default configuration without forcing or drag, the barotropic model's don't dependent on the previous prognostic state 
+    @test sum(to_vec(dprogn)[1]) ≈ 0 
+
+    # 
+    # dynamics_tendencies!
+    #
+
+    diagn_copy = deepcopy(diagn)
+    ddiag = one(diagn_copy)
+    ddiag_copy = deepcopy(ddiag)
+    progn_copy = deepcopy(progn)
+    dprogn = make_zero(progn)
+
+    autodiff(Reverse, SpeedyWeather.dynamics_tendencies!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(lf2), Duplicated(model, make_zero(model)))
+
+    function dynamics_tendencies(diagn, progn, lf, model)
+        diagn_new = deepcopy(diagn)
+        SpeedyWeather.dynamics_tendencies!(diagn_new, progn, lf, model)
+        return diagn_new
+    end 
+
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(9,1), x -> dynamics_tendencies(diagn_copy, x, lf2, model), ddiag_copy, progn_copy)
+
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
+
+    # in the default configuration without forcing or drag, the barotropic model's don't dependent on the previous prognostic state 
+    @test sum(to_vec(dprogn)[1]) ≈ 0 
+    
+    #
+    # Implicit correction 
+    #
+
+    diagn_copy = deepcopy(diagn)
+    ddiag = make_zero(diagn_copy)
+    ddiag_copy = deepcopy(ddiag)
+    progn_copy = deepcopy(progn)
+    dprogn = one(progn)
+
+    autodiff(Reverse, SpeedyWeather.implicit_correction!, Const, Duplicated(diagn, ddiag), Duplicated(model.implicit, make_zero(model.implicit)), Duplicated(progn, dprogn))
+
+    function implicit_correction(diagn, implicit, progn)
+        diagn_new = deepcopy(diagn)
+        SpeedyWeather.implicit_correction(diagn_new, implicit, progn)
+        return diagn_new
+    end 
+
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(9,1), x -> implicit_correction(diagn_copy, model.implicit, x), ddiag_copy, progn_copy)
+
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
+
+    #
+    # horizontal_diffusion!
+    #
+
+    lf1 = 1
+    diagn_copy = deepcopy(diagn)
+    ddiag = one(diagn_copy)
+    ddiag_copy = deepcopy(ddiag)
+
+    progn_copy = deepcopy(progn)
+    dprogn = make_zero(progn)
+
+    autodiff(Reverse, SpeedyWeather.horizontal_diffusion!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(model.horizontal_diffusion), Duplicated(model, make_zero(model)), Const(lf1))
+
+    # FD comparision not necessary, we have the exact values 
+    #function horizontal_diffusion(diagn, progn, diffusion, model, lf)
+    #    diagn_new = deepcopy(diagn)
+    #    SpeedyWeather.horizontal_diffusion!(diagn_new, progn, diffusion, model, lf)
+    #    return diagn_new
+    #end 
+
+    #fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> horizontal_diffusion(diagn_copy, x, model.horizontal_diffusion, model, lf1), ddiag_copy, progn_copy)
+    #@test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-2))
+
+    # ∂(progn)
+    # should be row-wise `model.horizontal_diffusion.impl .* model.horizontal_diffusion.expl`
+    # for all variables that are diffused 
+    diff_coefficient = model.horizontal_diffusion.impl .* model.horizontal_diffusion.expl
+    l_indices = [(1:l) for l=1:progn.vor[1].n]
+    for (i,il) in enumerate(l_indices)
+        @test all(real.(Matrix(dprogn.vor[lf1][:,1])[i, il]) .≈ diff_coefficient[i])
+    end 
+
+    # ∂(tend_old)
+    # should be row-wise `model.horizontal_diffusion.impl` 
+    for (i,il) in enumerate(l_indices)
+        @test all(real.(Matrix(ddiag.tendencies.vor_tend[:,1])[i, il]) .≈ model.horizontal_diffusion.impl[i])
+    end 
+
+    #
+    # Test the leapfrog 
+    # 
+
+    lf1 = 2 
+    lf2 = 2 
+
+    progn_copy = deepcopy(progn)
+    dprogn = one(progn_copy)
+    dprogn_copy = one(progn_copy)
+
+    tend = diagn.tendencies
+    tend_copy = deepcopy(tend)
+    dtend = make_zero(tend)
+    dmodel = make_zero(model)
+
+    autodiff(Reverse, SpeedyWeather.leapfrog!, Const, Duplicated(progn, dprogn), Duplicated(tend, dtend), Const(dt), Const(lf1), Const(model))
+
+    function leapfrog_step(progn_new::PrognosticVariables, progn::PrognosticVariables, tend, dt, lf, model)
+        copy!(progn_new, progn)
+        SpeedyWeather.leapfrog!(progn_new, tend, dt, lf, model)
+        return progn_new
+    end 
+
+    prog_new = zero(progn_copy)
+
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> leapfrog_step(prog_new, progn_copy, x, dt, lf1, model), dprogn_copy, tend_copy)
+
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dtend)[1],rtol=1e-5,atol=1e-5))
+
+    #
+    # transform!(diagn, progn, lf2, model)
+    #
+
+    fill!(diagn.tendencies, 0, PrimitiveWetModel)
+    diag_copy = deepcopy(diagn)
+    
+    ddiag = one(diagn)
+    ddiag_copy = deepcopy(ddiag)
+
+    progn_copy = deepcopy(progn)
+    dprogn = make_zero(progn)
+
+    autodiff(Reverse, SpeedyWeather.transform!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(lf2), Duplicated(model, make_zero(model)))
+
+    function transform_diagn(diag, progn, lf2, model)
+        diag_copy = deepcopy(diag)
+        transform!(diag_copy, progn, lf2, model)
+        return diag_copy
+    end 
+
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_diagn(diag_copy, x, lf2, model), ddiag_copy, progn_copy)
+    
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-3,atol=1e-3))
+end
+
+

From 36f916074e9ef3519269cb8e0bd1fbd187c921d3 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 31 Jan 2025 10:38:54 +0100
Subject: [PATCH 44/62] adjust progn zero to new land

---
 src/dynamics/prognostic_variables.jl | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl
index 2fbbc05f7..17bd28bbb 100644
--- a/src/dynamics/prognostic_variables.jl
+++ b/src/dynamics/prognostic_variables.jl
@@ -48,7 +48,7 @@ export PrognosticVariablesLand
     soil_temperature::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers)
 
     "Soil moisture, volume fraction [1]"
-    soil_moisture::GridVariable3D = zeros(GridVariable2D, nlat_half, nlayers)
+    soil_moisture::GridVariable3D = zeros(GridVariable3D, nlat_half, nlayers)
     
     "Snow depth [m]"
     snow_depth::GridVariable2D = zeros(GridVariable2D, nlat_half)
@@ -244,12 +244,12 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
 end
 
 function Base.zero(progn::PrognosticVariables{NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}) where {NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}
-    (; trunc, nlat_half, nlayers, nparticles) = progn
+    (; trunc, nlat_half, nlayers, nlayers_soil, nparticles) = progn
     
     # initialize regular progn variables 
     progn_new = PrognosticVariables{NF, ArrayType, nsteps,
         SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}(;
-            trunc, nlat_half, nlayers, nparticles,
+            trunc, nlat_half, nlayers, nlayers_soil, nparticles,
         )
 
     # add tracers with zero 
@@ -282,10 +282,9 @@ function Base.fill!(progn::PrognosticVariables{NF}, value::Number) where NF
     progn.ocean.evaporative_flux .= value_NF
 
     # land
-    progn.land.land_surface_temperature .= value_NF
+    progn.land.soil_temperature .= value_NF
     progn.land.snow_depth .= value_NF
-    progn.land.soil_moisture_layer1 .= value_NF
-    progn.land.soil_moisture_layer2 .= value_NF
+    progn.land.soil_moisture .= value_NF
     progn.land.sensible_heat_flux .= value_NF
     progn.land.evaporative_flux .= value_NF
 

From a968050312f635ab5320eca6c25189a7b8de27ea Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 31 Jan 2025 10:50:54 +0100
Subject: [PATCH 45/62] adjust zero for progn

---
 src/dynamics/prognostic_variables.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl
index 17bd28bbb..7261a3f22 100644
--- a/src/dynamics/prognostic_variables.jl
+++ b/src/dynamics/prognostic_variables.jl
@@ -248,7 +248,8 @@ function Base.zero(progn::PrognosticVariables{NF, ArrayType, nsteps, SpectralVar
     
     # initialize regular progn variables 
     progn_new = PrognosticVariables{NF, ArrayType, nsteps,
-        SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}(;
+        SpectralVariable2D, SpectralVariable3D, GridVariable2D, GridVariable3D, 
+        ParticleVector}(;
             trunc, nlat_half, nlayers, nlayers_soil, nparticles,
         )
 

From 0f472fd75ed166feb3c9fccacaec21819d64cfee Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 31 Jan 2025 11:12:12 +0100
Subject: [PATCH 46/62] actually fix it...

---
 src/dynamics/prognostic_variables.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl
index 7261a3f22..198cbfca4 100644
--- a/src/dynamics/prognostic_variables.jl
+++ b/src/dynamics/prognostic_variables.jl
@@ -243,7 +243,7 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
     return nothing
 end
 
-function Base.zero(progn::PrognosticVariables{NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}) where {NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D, GridVariable2D, ParticleVector}
+function Base.zero(progn::PrognosticVariables{NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D, GridVariable2D, GridVariable3D, ParticleVector}) where {NF, ArrayType, nsteps, SpectralVariable2D, SpectralVariable3D, GridVariable2D, GridVariable3D, ParticleVector}
     (; trunc, nlat_half, nlayers, nlayers_soil, nparticles) = progn
     
     # initialize regular progn variables 

From af14a22e949ffb82f3d1c953a6a02c448d786384 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 31 Jan 2025 16:25:53 +0100
Subject: [PATCH 47/62] replace NaNs in to_vec for FD

---
 ext/SpeedyWeatherFiniteDifferencesExt.jl | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/ext/SpeedyWeatherFiniteDifferencesExt.jl b/ext/SpeedyWeatherFiniteDifferencesExt.jl
index 477a54c93..371b36fa7 100644
--- a/ext/SpeedyWeatherFiniteDifferencesExt.jl
+++ b/ext/SpeedyWeatherFiniteDifferencesExt.jl
@@ -47,6 +47,7 @@ end
 
 # A version of the generic fallback from FiniteDifferences that excludes some of the fields 
 # that we don't want to be varied for our big data structures 
+# also replaces NaNs that are expected in land and ocean variables
 function FiniteDifferences.to_vec(x::T) where {T <: Union{PrognosticVariables, PrognosticVariablesOcean, PrognosticVariablesLand, DiagnosticVariables, Tendencies, GridVariables, DynamicsVariables, PhysicsVariables, ParticleVariables}}
 
     excluded_fields_pre, included_fields, excluded_fields_post = determine_included_fields(T)
@@ -59,6 +60,8 @@ function FiniteDifferences.to_vec(x::T) where {T <: Union{PrognosticVariables, P
     vals_excluded_post = map(name -> getfield(x, name), excluded_fields_post)
 
     v, vals_from_vec = to_vec(vals)
+    v = replace_NaN(x, v)
+
     function structtype_from_vec(v::Vector{<:Real})
         val_vecs = vals_from_vec(v)
         values = map((b, v) -> b(v), backs, val_vecs)
@@ -94,4 +97,21 @@ function determine_included_fields(T::Type)
     return excluded_fields_pre, included_fields, excluded_fields_post
 end 
 
+# in the ocean and land variables we have NaNs, FiniteDifferences can't deal with those, so we replace them
+function replace_NaN(x_type::T, vec) where {T <: Union{PrognosticVariablesOcean, PrognosticVariablesLand, PhysicsVariables}}
+    nan_indices = isnan.(vec)
+    vec[nan_indices] .= 0 
+    return vec
+end 
+
+# fallback, we really only want to replace the NaNs in ocean and land variables 
+replace_NaN(type, vec) = vec
+
+# By default FiniteDifferences doesn't include this, even though Integers can't be varied. 
+# there's an old GitHub issue and PR about this 
+function FiniteDifferences.to_vec(x::Integer)
+    Integer_from_vec(v) = x
+    return Bool[], Integer_from_vec
+end
+
 end 
\ No newline at end of file

From d05eb7d270fd16f6b94eb49701e510045b398c66 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 31 Jan 2025 16:26:23 +0100
Subject: [PATCH 48/62] WIP correctness tests

---
 test/differentiability/primitivewet.jl | 29 ++++++++++++++++++--------
 test/differentiability/timestepping.jl | 12 ++++++++---
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/test/differentiability/primitivewet.jl b/test/differentiability/primitivewet.jl
index 4721a7a2e..89386cfed 100644
--- a/test/differentiability/primitivewet.jl
+++ b/test/differentiability/primitivewet.jl
@@ -1,8 +1,12 @@
-### Experiments going a bit deeper into the timestepping of the barotropic model
+### Experiments going a bit deeper into the timestepping of the primtive wet model
+# this script / tests was mainly written for debugging, we might exclude it in future 
+# tests because it is quite maintanance heavy code 
 @testset "Differentiability: Primitive Wet Model Components" begin 
-    # T15 still yields somewhat sensible dynamics, that's why it's chosen here
-    spectral_grid = SpectralGrid(trunc=15, nlayers=3)          # define resolution
-    model = PrimitiveWetModel(; spectral_grid)   # construct model
+
+    spectral_grid = SpectralGrid(trunc=8, nlayers=1)          # define resolution
+
+    # FiniteDifferences struggles with the NaN when we have a land-sea mask, so we have to test on AquaPlanets 
+    model = PrimitiveWetModel(; spectral_grid)  # construct model
     simulation = initialize!(model)  
     initialize!(simulation)
     run!(simulation, period=Day(5)) # spin-up to get nonzero values for all fields
@@ -33,6 +37,7 @@
     dprogn = one(progn)
     ddiagn = one(diagn)
     ddiagn_copy = deepcopy(ddiagn)
+    diagn_copy = deepcopy(diagn)
 
     autodiff(Reverse, SpeedyWeather.parameterization_tendencies!, Const, Duplicated(diagn, ddiagn), Duplicated(progn, dprogn), Const(progn.clock.time), Duplicated(model, make_zero(model)))
 
@@ -42,8 +47,9 @@
         return diagn_new
     end 
 
-    # TODO: non-working at the moment -> NaN
     fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> parameterization_tendencies(diagn_copy, x, progn.clock.time, model), ddiagn_copy, progn_copy)
+    
+    # TO-DO this test is broken, they gradients don't line up
     #@test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
     
     #
@@ -65,12 +71,14 @@
     end 
 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> ocean_timestep(progn_copy, x, model), dprogn_copy, diagn_copy)
-    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
+
+    # also broken? 
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(ddiagn)[1],rtol=1e-4,atol=1e-1))
 
     #
     # land 
     # 
-    
+
     progn_copy = deepcopy(progn)
     dprogn = one(progn)
     ddiagn = make_zero(diagn)
@@ -86,7 +94,10 @@
     end 
 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> land_timestep(progn_copy, x, model), dprogn_copy, diagn_copy)
-    
+
+    # also broken currently
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(ddiagn)[1],rtol=1e-4,atol=1e-1))
+
     #####
     # DYNAMICS 
     lf2 = 2 
@@ -135,7 +146,7 @@
 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(9,1), x -> dynamics_tendencies(diagn_copy, x, lf2, model), ddiag_copy, progn_copy)
 
-    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
+    @test all(isapprox.(to_vec(fd_vjp)[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
 
     # in the default configuration without forcing or drag, the barotropic model's don't dependent on the previous prognostic state 
     @test sum(to_vec(dprogn)[1]) ≈ 0 
diff --git a/test/differentiability/timestepping.jl b/test/differentiability/timestepping.jl
index 66b5a2651..a27d31296 100644
--- a/test/differentiability/timestepping.jl
+++ b/test/differentiability/timestepping.jl
@@ -6,10 +6,11 @@
 
     for model_type in model_types 
         
-        nlayer = model_type == ShallowWaterModel ? 1 : 3
+        nlayer = model_type == ShallowWaterModel ? 1 : 1
 
-        spectral_grid = SpectralGrid(trunc=15, nlayers=nlayer)          # define resolution
-        model = model_type(; spectral_grid)   # construct model
+        # FiniteDifferences struggles with the NaN when we have a land-sea mask, so we have to test on AquaPlanets 
+        spectral_grid = SpectralGrid(trunc=8, nlayers=nlayer)          # define resolution
+        model = model_type(; spectral_grid, land_sea_mask = AquaPlanetMask(spectral_grid))   # construct model
         simulation = initialize!(model)  
         initialize!(simulation)
         run!(simulation, period=Day(3))
@@ -21,6 +22,9 @@
         progn = prognostic_variables
         diagn = diagnostic_variables
 
+        # TO-DO: The first time we execute this, the gradient is different. Why?
+        timestep_oop!(make_zero(progn), progn, diagn, dt, model)
+
         diagn_copy = deepcopy(diagn)
         progn_copy = deepcopy(progn)
 
@@ -39,6 +43,8 @@
         # FD comparison 
         dprogn_2 = one(progn) # seed 
 
+        # this takes a long time 
+        # with the FD comparision we have to go to quite low tolerences for the full time step 
         fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
         @test isapprox(to_vec(fd_vjp[1])[1], to_vec(d_progn)[1])

From 26354312e8b3072103bf71e9c2ddb9ffa664e0f6 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 31 Jan 2025 18:15:42 +0100
Subject: [PATCH 49/62] Implement dataids for array types (fixes broadcasting
 diff issue)

---
 .../lower_triangular_array.jl                 |  2 ++
 src/RingGrids/general.jl                      |  3 +++
 src/dynamics/prognostic_variables.jl          | 21 +++++++++----------
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/LowerTriangularMatrices/lower_triangular_array.jl b/src/LowerTriangularMatrices/lower_triangular_array.jl
index 6bab931be..c8ceb885c 100644
--- a/src/LowerTriangularMatrices/lower_triangular_array.jl
+++ b/src/LowerTriangularMatrices/lower_triangular_array.jl
@@ -123,6 +123,8 @@ function Base.array_summary(io::IO, L::LowerTriangularMatrix{T}, inds::Tuple{Var
     print(io, Base.dims2string(length.(inds)), ", $(mn[1])x$(mn[2]) LowerTriangularMatrix{$T}")
 end
 
+@inline Base.dataids(L::LowerTriangularArray) = Base.dataids(L.data)
+
 # CREATE INSTANCES (ZEROS, ONES, UNDEF)
 for f in (:zeros, :ones, :rand, :randn)
     @eval begin
diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl
index 1c1dd975f..e85ae912b 100644
--- a/src/RingGrids/general.jl
+++ b/src/RingGrids/general.jl
@@ -22,6 +22,9 @@ nonparametric_type(grid::AbstractGridArray) = nonparametric_type(typeof(grid))
 # also needed for other array types, defined in extensions
 nonparametric_type(::Type{<:Array}) = Array
 
+# needed for unalias 
+@inline Base.dataids(grid::AbstractGridArray) = Base.dataids(grid.data)
+
 """$(TYPEDSIGNATURES) Full grid array type for `grid`. Always returns the N-dimensional `*Array`
 not the two-dimensional (`N=1`) `*Grid`. For reduced grids the corresponding full grid that
 share the same latitudes."""
diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl
index 198cbfca4..d1d058263 100644
--- a/src/dynamics/prognostic_variables.jl
+++ b/src/dynamics/prognostic_variables.jl
@@ -205,19 +205,18 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
         progn_new.pres[i] .= progn_old.pres[i]
     end
 
-    # TO-DO: currently there are some problems with Enzyme and RingGrids broadcasting
     # ocean
-    progn_new.ocean.sea_surface_temperature.data .= progn_old.ocean.sea_surface_temperature.data
-    progn_new.ocean.sea_ice_concentration.data .= progn_old.ocean.sea_ice_concentration.data
-    progn_new.ocean.sensible_heat_flux.data .= progn_old.ocean.sensible_heat_flux.data
-    progn_new.ocean.evaporative_flux.data .= progn_old.ocean.evaporative_flux.data
+    progn_new.ocean.sea_surface_temperature .= progn_old.ocean.sea_surface_temperature
+    progn_new.ocean.sea_ice_concentration .= progn_old.ocean.sea_ice_concentration
+    progn_new.ocean.sensible_heat_flux .= progn_old.ocean.sensible_heat_flux
+    progn_new.ocean.evaporative_flux .= progn_old.ocean.evaporative_flux
  
     # land
-    progn_new.land.soil_temperature.data .= progn_old.land.soil_temperature.data
-    progn_new.land.snow_depth.data .= progn_old.land.snow_depth.data
-    progn_new.land.soil_moisture.data .= progn_old.land.soil_moisture.data
-    progn_new.land.sensible_heat_flux.data .= progn_old.land.sensible_heat_flux.data
-    progn_new.land.evaporative_flux.data .= progn_old.land.evaporative_flux.data
+    progn_new.land.soil_temperature .= progn_old.land.soil_temperature
+    progn_new.land.snow_depth .= progn_old.land.snow_depth
+    progn_new.land.soil_moisture .= progn_old.land.soil_moisture
+    progn_new.land.sensible_heat_flux .= progn_old.land.sensible_heat_flux
+    progn_new.land.evaporative_flux .= progn_old.land.evaporative_flux
 
     # copy over tracers
     for (key, value) in progn_old.tracers
@@ -235,7 +234,7 @@ function Base.copy!(progn_new::PrognosticVariables, progn_old::PrognosticVariabl
         progn_new.particles .= progn_old.particles
     end
 
-    progn_new.random_pattern.data .= progn_old.random_pattern.data
+    progn_new.random_pattern .= progn_old.random_pattern.data
 
     copy!(progn_new.clock, progn_old.clock)
     progn_new.scale[] = progn_old.scale[]

From 337567425d439110aeae815e46e0c5afc417a333 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 31 Jan 2025 19:25:37 +0100
Subject: [PATCH 50/62] adjust to now functioning ringgrids broadcast

---
 src/dynamics/tendencies.jl                     | 14 +++++++-------
 test/differentiability/timestepping.jl         |  2 +-
 test/transforms/spectral_transform_ad_rules.jl |  8 +++++++-
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/dynamics/tendencies.jl b/src/dynamics/tendencies.jl
index 4ee807b34..646272a07 100644
--- a/src/dynamics/tendencies.jl
+++ b/src/dynamics/tendencies.jl
@@ -853,10 +853,10 @@ function SpeedyTransforms.transform!(
 
     # retain previous time step for vertical advection and some parameterizations
     if initialize == false              # only store prev after initial step
-        @. temp_grid_prev.data = temp_grid.data   # this is temperature anomaly wrt to implicit reference profile!
-        @. humid_grid_prev.data = humid_grid.data
-        @. u_grid_prev.data = u_grid.data
-        @. v_grid_prev.data = v_grid.data
+        @. temp_grid_prev = temp_grid   # this is temperature anomaly wrt to implicit reference profile!
+        @. humid_grid_prev = humid_grid
+        @. u_grid_prev = u_grid
+        @. v_grid_prev = v_grid
         # TODO: above broadcastings are done with the .data because there's currently a problem with Enzyme and our RingGrids broadcasting
 
         for (name, tracer) in model.tracers
@@ -901,9 +901,9 @@ function SpeedyTransforms.transform!(
             end
         end
 
-        @. humid_grid_prev.data = humid_grid.data
-        @. u_grid_prev.data = u_grid.data
-        @. v_grid_prev.data = v_grid.data
+        @. humid_grid_prev = humid_grid
+        @. u_grid_prev = u_grid
+        @. v_grid_prev = v_grid
 
         for (name, tracer) in model.tracers
             if tracer.active
diff --git a/test/differentiability/timestepping.jl b/test/differentiability/timestepping.jl
index a27d31296..4c6183ade 100644
--- a/test/differentiability/timestepping.jl
+++ b/test/differentiability/timestepping.jl
@@ -10,7 +10,7 @@
 
         # FiniteDifferences struggles with the NaN when we have a land-sea mask, so we have to test on AquaPlanets 
         spectral_grid = SpectralGrid(trunc=8, nlayers=nlayer)          # define resolution
-        model = model_type(; spectral_grid, land_sea_mask = AquaPlanetMask(spectral_grid))   # construct model
+        model = model_type(; spectral_grid)   # construct model
         simulation = initialize!(model)  
         initialize!(simulation)
         run!(simulation, period=Day(3))
diff --git a/test/transforms/spectral_transform_ad_rules.jl b/test/transforms/spectral_transform_ad_rules.jl
index c9030daca..1f1e042df 100644
--- a/test/transforms/spectral_transform_ad_rules.jl
+++ b/test/transforms/spectral_transform_ad_rules.jl
@@ -86,6 +86,12 @@ end
     progn_new = zero(progn)
     dprogn_new = one(progn) # seed 
 
-    autodiff(Reverse, SpeedyWeather.timestep!, Const, Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
+    function timestep_oop!(progn_new::PrognosticVariables, progn_old::PrognosticVariables, diagn, dt, model, lf1=2, lf2=2)
+        copy!(progn_new, progn_old)
+        SpeedyWeather.timestep!(progn_new, diagn, dt, model, lf1, lf2)
+        return nothing
+    end 
+
+    autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
     @test sum(to_vec(d_progn)[1]) != 0
 end 
\ No newline at end of file

From 796b331e6a2b540316668b3e126dbc286850553d Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Fri, 31 Jan 2025 19:28:03 +0100
Subject: [PATCH 51/62] reduce resolution of CI test

---
 test/transforms/spectral_transform_ad_rules.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/transforms/spectral_transform_ad_rules.jl b/test/transforms/spectral_transform_ad_rules.jl
index 1f1e042df..2e2fa53dd 100644
--- a/test/transforms/spectral_transform_ad_rules.jl
+++ b/test/transforms/spectral_transform_ad_rules.jl
@@ -63,7 +63,7 @@ end
     # in a seperate test set. But we do want to ensure in the regular CI that 
     # we don't commit some kind of problem for the Enzyme differentiability
     # so, we test here if we get a non-zero gradient from the timestepping.  
-    spectral_grid = SpectralGrid(trunc=15, nlayers=3)          # define resolution
+    spectral_grid = SpectralGrid(trunc=8, nlayers=1)          # define resolution
     model = PrimitiveWetModel(; spectral_grid)   # construct model
     simulation = initialize!(model)  
     initialize!(simulation)

From 6a7767e4c4569e9465fa4196f58c6b44b53712b1 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Tue, 4 Feb 2025 19:52:05 +0100
Subject: [PATCH 52/62] copy all structs for FD

---
 test/differentiability/timestep_utils.jl | 13 ++++++++++---
 test/differentiability/timestepping.jl   | 23 ++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/test/differentiability/timestep_utils.jl b/test/differentiability/timestep_utils.jl
index 090a51940..043d237a4 100644
--- a/test/differentiability/timestep_utils.jl
+++ b/test/differentiability/timestep_utils.jl
@@ -4,8 +4,15 @@ function timestep_oop!(progn_new::PrognosticVariables, progn_old::PrognosticVari
     return nothing
 end 
 
+# for FiniteDifferences.jl, we need to copy all inputs that are mutated 
+# because this function is called many times by FiniteDifferences
 function timestep_oop(progn, diagn, dt, model, lf1=2, lf2=2)
-    progn_new = zero(progn)
-    timestep_oop!(progn_new, progn, diagn, dt, model, lf1, lf2)
-    return progn_new
+
+    progn_copy = deepcopy(progn)
+    diagn_copy = deepcopy(diagn)
+
+    model_copy = deepcopy(model) # just to be save, as we have some temporary memory in there as well
+
+    SpeedyWeather.timestep!(progn_copy, diagn_copy, dt, model_copy, lf1, lf2)
+    return progn_copy
 end 
diff --git a/test/differentiability/timestepping.jl b/test/differentiability/timestepping.jl
index 4c6183ade..c719c2ae4 100644
--- a/test/differentiability/timestepping.jl
+++ b/test/differentiability/timestepping.jl
@@ -28,14 +28,19 @@
         diagn_copy = deepcopy(diagn)
         progn_copy = deepcopy(progn)
 
+        diagn_copy_2 = deepcopy(diagn) # for a later experiment a second copy
+        progn_copy_2 = deepcopy(progn)
+
         d_progn = zero(progn)
         d_diag = make_zero(diagn)
 
         progn_new = zero(progn)
         dprogn_new = one(progn) # seed 
 
+        dmodel = make_zero(model)
+
         # test that we can differentiate wrt an IC 
-        autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, make_zero(model)))
+        autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, dmodel))
 
         # nonzero gradient
         @test sum(to_vec(d_progn)[1]) != 0
@@ -48,5 +53,21 @@
         fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> timestep_oop(x, diagn_copy, dt, model), dprogn_2, progn_copy)
 
         @test isapprox(to_vec(fd_vjp[1])[1], to_vec(d_progn)[1])
+
+        # wrt a system parameter. Let's check for example gravity 
+
+        # function wrapper for FiniteDifferences
+        function timestep_wrt_gravity(g)
+            model_new = deepcopy(model)
+            model_new.planet.gravity = g 
+            println(g)
+            timestep_oop(progn_copy_2, diagn_copy_2, dt, model_new)
+        end   
+        
+        dprogn_2 = one(progn) # seed
+
+        fd_vjp = FiniteDifferences.j′vp(central_fdm(11,1), timestep_wrt_gravity, dprogn_2, model.planet.gravity)
+        # this doesn't line up, currently, what's wrong? 
+
     end 
 end 
\ No newline at end of file

From 06de07df05b61a4006cd42aeda21660ea336eb04 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Tue, 4 Feb 2025 19:58:51 +0100
Subject: [PATCH 53/62] parameter test in timestepping as well

---
 test/differentiability/timestepping.jl | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/test/differentiability/timestepping.jl b/test/differentiability/timestepping.jl
index c719c2ae4..7af87da1a 100644
--- a/test/differentiability/timestepping.jl
+++ b/test/differentiability/timestepping.jl
@@ -58,16 +58,17 @@
 
         # function wrapper for FiniteDifferences
         function timestep_wrt_gravity(g)
+            progn_fd = deepcopy(progn_copy_2)
+            diagn_fd = deepcopy(diagn_copy_2)
             model_new = deepcopy(model)
             model_new.planet.gravity = g 
-            println(g)
-            timestep_oop(progn_copy_2, diagn_copy_2, dt, model_new)
+            timestep_oop(progn_fd, diagn_fd, dt, model_new)
         end   
         
         dprogn_2 = one(progn) # seed
 
         fd_vjp = FiniteDifferences.j′vp(central_fdm(11,1), timestep_wrt_gravity, dprogn_2, model.planet.gravity)
-        # this doesn't line up, currently, what's wrong? 
 
+        @test isapprox(dmodel.planet.gravity, fd_vjp[1],rtol=1e-1)
     end 
 end 
\ No newline at end of file

From 16dc4122249dca2314183fe46599602c81b49dbb Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 5 Feb 2025 08:46:13 +0100
Subject: [PATCH 54/62] doc update

---
 docs/src/differentiability.md          | 53 +++++++++++++++++++++++++-
 test/differentiability/primitivewet.jl |  5 ++-
 test/differentiability/runtests.jl     |  1 +
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/docs/src/differentiability.md b/docs/src/differentiability.md
index ad55fc03d..d07ce126a 100644
--- a/docs/src/differentiability.md
+++ b/docs/src/differentiability.md
@@ -3,6 +3,55 @@
 SpeedyWeather.jl is written with differentiability in mind. This means that our model is differentiable by automatic differentiation (AD). If you are interested in machine learning (ML), this means that you can integrate our model directly into your ML models without the need to train your ANNs seperately first. For atmospheric modellers this means that you get an adjoint model for free which is always generated automatically, so that we don't need to maintain it seperatly. So, you can calibrate SpeedyWeather.jl in a fully automatic, data-driven way. 
 
 !!! Work in progress
-    The differentiability of SpeedyWeather.jl is still work in progress and some parts of this documentation might be not be always updated to the latest state. We will extend this documentation over time. Don't hesitate to contact us via GitHub issues or mail when you have questions. 
+    The differentiability of SpeedyWeather.jl is still work in progress and some parts of this documentation might be not be always updated to the latest state. We will extend this documentation over time. Don't hesitate to contact us via GitHub issues or mail when you have questions or want to colloborate.
+
+For the differentiability of our model we rely on [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl). If you've used Enzyme before, just go ahead and try to differentiate the model! It should work. We have checked the correctness of the gradients extensively against a finite differences differentiation with [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl/). In the following we present a simple example how we can take the gradient of a single timestep of the primitive equation model with respect to one of the model parameter. 
+
+!!! Enzyme with Julia 1.11
+    Currently there are still some issues with Enzyme in Julia 1.11, we recommend to use Julia 1.10 for the following
+
+First we initialize the model as usual: 
+
+```@example autodiff
+using SpeedyWeather, Enzyme 
+
+spectral_grid = SpectralGrid(trunc=23, nlayers=3)           
+model = PrimitiveWetModel(; spectral_grid) 
+simulation = initialize!(model)  
+initialize!(simulation)
+run!(simulation, period=Day(10)) # spin-up the model a bit
+```
+
+Then, we get all variables we need from our `simulation`
+
+```@example autodiff
+(; prognostic_variables, diagnostic_variables, model) = simulation
+(; Δt, Δt_millisec) = model.time_stepping
+dt = 2Δt
+
+progn = prognostic_variables
+diagn = diagnostic_variables
+```
+
+Next, we will prepare to use Enzyme. Enzyme saves the gradient information in a shadow of the original input. For the inputs this shadow is initialized zero, whereas for the output the shadow is used as the seed of the AD. In other words, as we are doing reverse-mode AD, the shadow of the output is the value that is backpropageted by the reverse-mode AD. Ok, let's initialize everything: 
+
+```@example autodiff 
+dprogn = one(progn) # shadow for the progn values 
+ddiagn = make_zero(diagn) # shadow for the diagn values 
+dmodel = make_zero(model) # here, we'll accumulate all parameter derivatives 
+```
+
+Then, we can already do the differentiation with Enzyme
+
+```@example autodiff 
+autodiff(Reverse, SpeedyWeather.timestep!, Const, Duplicated(progn, dprogn), Duplicated(diagn, ddiagn), Const(dt), Duplicated(model, dmodel))
+```
+
+The derivitaves are accumulated in the `dmodel` shadow. So, if we e.g. want to know the derivative with respect to the gravity constant, we just have to inspect: 
+
+```@example autodiff 
+dmodel.planet.gravity 
+```
+
+Doing a full sensitivity analysis through a long integration is computationally much more demanding, and is something that we are currently working on. 
 
-For the differentiability of our model we rely on [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl). If you've used Enzyme before, just go ahead and try to differnentiate the model! It should work. We have checked the correctness of the gradients extensively against a finite differences differentiation with [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl/). In the following we present a simple example how we can take the gradient of a single timestep of the primitive equation model with respect to one of the model parameter, so with ``\mathbf{P}_t = (\zeta, \mathcal{D}, p_s, T, q)`` as the state vector of all prognostic variables at time step
\ No newline at end of file
diff --git a/test/differentiability/primitivewet.jl b/test/differentiability/primitivewet.jl
index 89386cfed..14f9108a2 100644
--- a/test/differentiability/primitivewet.jl
+++ b/test/differentiability/primitivewet.jl
@@ -1,5 +1,5 @@
 ### Experiments going a bit deeper into the timestepping of the primtive wet model
-# this script / tests was mainly written for debugging, we might exclude it in future 
+# this script / these tests were mainly written for debugging, we might exclude it in future 
 # tests because it is quite maintanance heavy code 
 @testset "Differentiability: Primitive Wet Model Components" begin 
 
@@ -33,6 +33,7 @@
     #
     # model physics 
     # 
+
     progn_copy = deepcopy(progn)
     dprogn = one(progn)
     ddiagn = one(diagn)
@@ -43,7 +44,7 @@
 
     function parameterization_tendencies(diagn, progn, time, model)
         diagn_new = deepcopy(diagn)
-        SpeedyWeather.parameterization_tendencies!(diagn_new, progn, time, model)
+        SpeedyWeather.parameterization_tendencies!(diagn_new, deepcopy(progn), time, deepcopy(model))
         return diagn_new
     end 
 
diff --git a/test/differentiability/runtests.jl b/test/differentiability/runtests.jl
index 13881c875..bc5a292d9 100644
--- a/test/differentiability/runtests.jl
+++ b/test/differentiability/runtests.jl
@@ -13,4 +13,5 @@ include("timestep_utils.jl")
 # TESTS 
 include("speedy_transforms.jl")
 include("barotropic.jl")
+#include("primitivewet.jl")
 include("timestepping.jl")
\ No newline at end of file

From cff60e1da54cbda7346afb18ed2ce2789d94a274 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 5 Feb 2025 08:57:39 +0100
Subject: [PATCH 55/62] remove to-do note

---
 src/dynamics/tendencies.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/dynamics/tendencies.jl b/src/dynamics/tendencies.jl
index 646272a07..5ff45fdf3 100644
--- a/src/dynamics/tendencies.jl
+++ b/src/dynamics/tendencies.jl
@@ -857,7 +857,6 @@ function SpeedyTransforms.transform!(
         @. humid_grid_prev = humid_grid
         @. u_grid_prev = u_grid
         @. v_grid_prev = v_grid
-        # TODO: above broadcastings are done with the .data because there's currently a problem with Enzyme and our RingGrids broadcasting
 
         for (name, tracer) in model.tracers
             if tracer.active

From 7bda523a5ef214dc7330f059e3f8f25cb4454ea6 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 5 Feb 2025 09:02:32 +0100
Subject: [PATCH 56/62] add enzyme to docs env

---
 docs/Project.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/Project.toml b/docs/Project.toml
index 4d8ff194c..3ac8cb0b1 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -1,6 +1,7 @@
 [deps]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"

From 42ce046101976b17bfc59f7819c2ab388bb93738 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Wed, 5 Feb 2025 10:38:31 +0100
Subject: [PATCH 57/62] remove Enzyme dep from docs again

---
 docs/Project.toml             |  1 -
 docs/src/differentiability.md | 11 +++++------
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/docs/Project.toml b/docs/Project.toml
index 3ac8cb0b1..4d8ff194c 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -1,7 +1,6 @@
 [deps]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
-Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
diff --git a/docs/src/differentiability.md b/docs/src/differentiability.md
index d07ce126a..b69a25790 100644
--- a/docs/src/differentiability.md
+++ b/docs/src/differentiability.md
@@ -12,7 +12,7 @@ For the differentiability of our model we rely on [Enzyme.jl](https://github.com
 
 First we initialize the model as usual: 
 
-```@example autodiff
+```julia
 using SpeedyWeather, Enzyme 
 
 spectral_grid = SpectralGrid(trunc=23, nlayers=3)           
@@ -24,7 +24,7 @@ run!(simulation, period=Day(10)) # spin-up the model a bit
 
 Then, we get all variables we need from our `simulation`
 
-```@example autodiff
+```julia
 (; prognostic_variables, diagnostic_variables, model) = simulation
 (; Δt, Δt_millisec) = model.time_stepping
 dt = 2Δt
@@ -35,7 +35,7 @@ diagn = diagnostic_variables
 
 Next, we will prepare to use Enzyme. Enzyme saves the gradient information in a shadow of the original input. For the inputs this shadow is initialized zero, whereas for the output the shadow is used as the seed of the AD. In other words, as we are doing reverse-mode AD, the shadow of the output is the value that is backpropageted by the reverse-mode AD. Ok, let's initialize everything: 
 
-```@example autodiff 
+```julia
 dprogn = one(progn) # shadow for the progn values 
 ddiagn = make_zero(diagn) # shadow for the diagn values 
 dmodel = make_zero(model) # here, we'll accumulate all parameter derivatives 
@@ -43,15 +43,14 @@ dmodel = make_zero(model) # here, we'll accumulate all parameter derivatives
 
 Then, we can already do the differentiation with Enzyme
 
-```@example autodiff 
+```julia
 autodiff(Reverse, SpeedyWeather.timestep!, Const, Duplicated(progn, dprogn), Duplicated(diagn, ddiagn), Const(dt), Duplicated(model, dmodel))
 ```
 
 The derivitaves are accumulated in the `dmodel` shadow. So, if we e.g. want to know the derivative with respect to the gravity constant, we just have to inspect: 
 
-```@example autodiff 
+```julia 
 dmodel.planet.gravity 
 ```
 
 Doing a full sensitivity analysis through a long integration is computationally much more demanding, and is something that we are currently working on. 
-

From f82883a64c3ce365aeb20555198c66e282a1b6fc Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Thu, 13 Feb 2025 19:53:27 +0100
Subject: [PATCH 58/62] full diff test only in Julia 1.10

---
 .../transforms/spectral_transform_ad_rules.jl | 64 ++++++++++---------
 1 file changed, 34 insertions(+), 30 deletions(-)

diff --git a/test/transforms/spectral_transform_ad_rules.jl b/test/transforms/spectral_transform_ad_rules.jl
index 2e2fa53dd..732299b98 100644
--- a/test/transforms/spectral_transform_ad_rules.jl
+++ b/test/transforms/spectral_transform_ad_rules.jl
@@ -58,40 +58,44 @@ end
     end
 end 
 
-@testset "Complete Differentiability" begin 
-    # We do extensive correctness checks and tests on the differentiability 
-    # in a seperate test set. But we do want to ensure in the regular CI that 
-    # we don't commit some kind of problem for the Enzyme differentiability
-    # so, we test here if we get a non-zero gradient from the timestepping.  
-    spectral_grid = SpectralGrid(trunc=8, nlayers=1)          # define resolution
-    model = PrimitiveWetModel(; spectral_grid)   # construct model
-    simulation = initialize!(model)  
-    initialize!(simulation)
-    run!(simulation, period=Day(1))
-    
-    (; prognostic_variables, diagnostic_variables, model) = simulation
-    (; Δt, Δt_millisec) = model.time_stepping
-    dt = 2Δt
+# Enzyme and Julia 1.11 still has some problems, and the test below is broken
+# in Julia 1.11
+if VERSION >= v"1.11.0"
+    @testset "Complete Differentiability" begin 
+        # We do extensive correctness checks and tests on the differentiability 
+        # in a seperate test set. But we do want to ensure in the regular CI that 
+        # we don't commit some kind of problem for the Enzyme differentiability
+        # so, we test here if we get a non-zero gradient from the timestepping.  
+        spectral_grid = SpectralGrid(trunc=8, nlayers=1)          # define resolution
+        model = PrimitiveWetModel(; spectral_grid)   # construct model
+        simulation = initialize!(model)  
+        initialize!(simulation)
+        run!(simulation, period=Day(1))
+        
+        (; prognostic_variables, diagnostic_variables, model) = simulation
+        (; Δt, Δt_millisec) = model.time_stepping
+        dt = 2Δt
 
-    progn = prognostic_variables
-    diagn = diagnostic_variables
+        progn = prognostic_variables
+        diagn = diagnostic_variables
 
-    diagn_copy = deepcopy(diagn)
-    progn_copy = deepcopy(progn)
+        diagn_copy = deepcopy(diagn)
+        progn_copy = deepcopy(progn)
 
-    d_progn = zero(progn)
-    d_diag = make_zero(diagn)
-    d_model = make_zero(model)
+        d_progn = zero(progn)
+        d_diag = make_zero(diagn)
+        d_model = make_zero(model)
 
-    progn_new = zero(progn)
-    dprogn_new = one(progn) # seed 
+        progn_new = zero(progn)
+        dprogn_new = one(progn) # seed 
 
-    function timestep_oop!(progn_new::PrognosticVariables, progn_old::PrognosticVariables, diagn, dt, model, lf1=2, lf2=2)
-        copy!(progn_new, progn_old)
-        SpeedyWeather.timestep!(progn_new, diagn, dt, model, lf1, lf2)
-        return nothing
-    end 
+        function timestep_oop!(progn_new::PrognosticVariables, progn_old::PrognosticVariables, diagn, dt, model, lf1=2, lf2=2)
+            copy!(progn_new, progn_old)
+            SpeedyWeather.timestep!(progn_new, diagn, dt, model, lf1, lf2)
+            return nothing
+        end 
 
-    autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
-    @test sum(to_vec(d_progn)[1]) != 0
+        autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
+        @test sum(to_vec(d_progn)[1]) != 0
+    end 
 end 
\ No newline at end of file

From 926d699093db8e0301545c93ad4c903432244ecf Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Thu, 13 Feb 2025 20:20:29 +0100
Subject: [PATCH 59/62] actually restrict Julia version correclty ...

---
 test/transforms/spectral_transform_ad_rules.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/transforms/spectral_transform_ad_rules.jl b/test/transforms/spectral_transform_ad_rules.jl
index 732299b98..5bc74b621 100644
--- a/test/transforms/spectral_transform_ad_rules.jl
+++ b/test/transforms/spectral_transform_ad_rules.jl
@@ -60,7 +60,7 @@ end
 
 # Enzyme and Julia 1.11 still has some problems, and the test below is broken
 # in Julia 1.11
-if VERSION >= v"1.11.0"
+if VERSION <= v"1.11.0"
     @testset "Complete Differentiability" begin 
         # We do extensive correctness checks and tests on the differentiability 
         # in a seperate test set. But we do want to ensure in the regular CI that 

From 57b50cfc06868b190040a901d9d6ce6caab608af Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Mon, 17 Feb 2025 15:46:07 +0100
Subject: [PATCH 60/62] exclude for faulty land constructor

---
 src/physics/land/land.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/physics/land/land.jl b/src/physics/land/land.jl
index 3b5e5b3ba..2254ba68a 100644
--- a/src/physics/land/land.jl
+++ b/src/physics/land/land.jl
@@ -15,14 +15,14 @@ export LandModel
 end
 
 # default constructor
-function LandModel(SG::SpectralGrid)
+#=function LandModel(SG::SpectralGrid)
     return LandModel(
         SeasonalLandTemperature(SG),
         SeasonalSoilMoisture(SG),
         VegetationClimatology(SG),
         NoRivers(SG),
     )
-end
+end =#
 
 function LandModel(
     SG::SpectralGrid;

From f1cac4e36303bfbe562ad478cbeb91e104e3ef9a Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Mon, 17 Feb 2025 17:35:31 +0100
Subject: [PATCH 61/62] minor work on primtive wet detailed checks

---
 test/differentiability/primitivewet.jl | 111 ++-----------------------
 1 file changed, 7 insertions(+), 104 deletions(-)

diff --git a/test/differentiability/primitivewet.jl b/test/differentiability/primitivewet.jl
index 14f9108a2..d48b5bbbf 100644
--- a/test/differentiability/primitivewet.jl
+++ b/test/differentiability/primitivewet.jl
@@ -67,14 +67,13 @@
 
     function ocean_timestep(progn, diagn, model)
         progn_new = deepcopy(progn)
-        SpeedyWeather.ocean_timestep!(progn_new, diagn, model)
+        SpeedyWeather.ocean_timestep!(progn_new, deepcopy(diagn), deepcopy(model))
         return progn_new
     end 
 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> ocean_timestep(progn_copy, x, model), dprogn_copy, diagn_copy)
 
-    # also broken? 
-    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(ddiagn)[1],rtol=1e-4,atol=1e-1))
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(ddiagn)[1],rtol=1e-4,atol=1e-2))
 
     #
     # land 
@@ -90,43 +89,18 @@
 
     function land_timestep(progn, diagn, model)
         progn_new = deepcopy(progn)
-        SpeedyWeather.ocean_timestep!(progn_new, diagn, model)
+        SpeedyWeather.ocean_timestep!(progn_new, deepcopy(diagn), deepcopy(model))
         return progn_new
     end 
 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> land_timestep(progn_copy, x, model), dprogn_copy, diagn_copy)
 
-    # also broken currently
-    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(ddiagn)[1],rtol=1e-4,atol=1e-1))
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(ddiagn)[1],rtol=1e-4,atol=1e-2))
 
     #####
     # DYNAMICS 
     lf2 = 2 
-    #
-    # drag! 
-    #
-
-    diagn_copy = deepcopy(diagn)
-    ddiag = one(diagn_copy)
-    ddiag_copy = deepcopy(ddiag)
-    progn_copy = deepcopy(progn)
-    dprogn = make_zero(progn)
-
-    autodiff(Reverse, SpeedyWeather.drag!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(lf2), Duplicated(model, make_zero(model)))
-
-    function drag(diagn, progn, lf, model)
-        diagn_new = deepcopy(diagn)
-        SpeedyWeather.drag!(diagn_new, progn, lf, model)
-        return diagn_new
-    end 
-
-    fd_vjp = FiniteDifferences.j′vp(central_fdm(9,1), x -> drag(diagn_copy, x, lf2, model), ddiag_copy, progn_copy)
-
-    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
-
-    # in the default configuration without forcing or drag, the barotropic model's don't dependent on the previous prognostic state 
-    @test sum(to_vec(dprogn)[1]) ≈ 0 
-
+   
     # 
     # dynamics_tendencies!
     #
@@ -141,16 +115,14 @@
 
     function dynamics_tendencies(diagn, progn, lf, model)
         diagn_new = deepcopy(diagn)
-        SpeedyWeather.dynamics_tendencies!(diagn_new, progn, lf, model)
+        SpeedyWeather.dynamics_tendencies!(diagn_new, deepcopy(progn), lf, deepcopy(model))
         return diagn_new
     end 
 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(9,1), x -> dynamics_tendencies(diagn_copy, x, lf2, model), ddiag_copy, progn_copy)
 
+    # there are some NaNs in the FD, that's why this test is currently broken
     @test all(isapprox.(to_vec(fd_vjp)[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
-
-    # in the default configuration without forcing or drag, the barotropic model's don't dependent on the previous prognostic state 
-    @test sum(to_vec(dprogn)[1]) ≈ 0 
     
     #
     # Implicit correction 
@@ -174,75 +146,6 @@
 
     @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
 
-    #
-    # horizontal_diffusion!
-    #
-
-    lf1 = 1
-    diagn_copy = deepcopy(diagn)
-    ddiag = one(diagn_copy)
-    ddiag_copy = deepcopy(ddiag)
-
-    progn_copy = deepcopy(progn)
-    dprogn = make_zero(progn)
-
-    autodiff(Reverse, SpeedyWeather.horizontal_diffusion!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(model.horizontal_diffusion), Duplicated(model, make_zero(model)), Const(lf1))
-
-    # FD comparision not necessary, we have the exact values 
-    #function horizontal_diffusion(diagn, progn, diffusion, model, lf)
-    #    diagn_new = deepcopy(diagn)
-    #    SpeedyWeather.horizontal_diffusion!(diagn_new, progn, diffusion, model, lf)
-    #    return diagn_new
-    #end 
-
-    #fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> horizontal_diffusion(diagn_copy, x, model.horizontal_diffusion, model, lf1), ddiag_copy, progn_copy)
-    #@test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-2))
-
-    # ∂(progn)
-    # should be row-wise `model.horizontal_diffusion.impl .* model.horizontal_diffusion.expl`
-    # for all variables that are diffused 
-    diff_coefficient = model.horizontal_diffusion.impl .* model.horizontal_diffusion.expl
-    l_indices = [(1:l) for l=1:progn.vor[1].n]
-    for (i,il) in enumerate(l_indices)
-        @test all(real.(Matrix(dprogn.vor[lf1][:,1])[i, il]) .≈ diff_coefficient[i])
-    end 
-
-    # ∂(tend_old)
-    # should be row-wise `model.horizontal_diffusion.impl` 
-    for (i,il) in enumerate(l_indices)
-        @test all(real.(Matrix(ddiag.tendencies.vor_tend[:,1])[i, il]) .≈ model.horizontal_diffusion.impl[i])
-    end 
-
-    #
-    # Test the leapfrog 
-    # 
-
-    lf1 = 2 
-    lf2 = 2 
-
-    progn_copy = deepcopy(progn)
-    dprogn = one(progn_copy)
-    dprogn_copy = one(progn_copy)
-
-    tend = diagn.tendencies
-    tend_copy = deepcopy(tend)
-    dtend = make_zero(tend)
-    dmodel = make_zero(model)
-
-    autodiff(Reverse, SpeedyWeather.leapfrog!, Const, Duplicated(progn, dprogn), Duplicated(tend, dtend), Const(dt), Const(lf1), Const(model))
-
-    function leapfrog_step(progn_new::PrognosticVariables, progn::PrognosticVariables, tend, dt, lf, model)
-        copy!(progn_new, progn)
-        SpeedyWeather.leapfrog!(progn_new, tend, dt, lf, model)
-        return progn_new
-    end 
-
-    prog_new = zero(progn_copy)
-
-    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> leapfrog_step(prog_new, progn_copy, x, dt, lf1, model), dprogn_copy, tend_copy)
-
-    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dtend)[1],rtol=1e-5,atol=1e-5))
-
     #
     # transform!(diagn, progn, lf2, model)
     #

From dee92e33868152ff051ba11addace708fc29cd91 Mon Sep 17 00:00:00 2001
From: Max <maximilian.gelbrecht@posteo.de>
Date: Tue, 18 Feb 2025 15:48:42 +0100
Subject: [PATCH 62/62] mark broken test in 1.11

---
 test/differentiability/primitivewet.jl        | 20 +++++++++----------
 .../transforms/spectral_transform_ad_rules.jl |  4 ++++
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/test/differentiability/primitivewet.jl b/test/differentiability/primitivewet.jl
index d48b5bbbf..ec9482246 100644
--- a/test/differentiability/primitivewet.jl
+++ b/test/differentiability/primitivewet.jl
@@ -5,7 +5,6 @@
 
     spectral_grid = SpectralGrid(trunc=8, nlayers=1)          # define resolution
 
-    # FiniteDifferences struggles with the NaN when we have a land-sea mask, so we have to test on AquaPlanets 
     model = PrimitiveWetModel(; spectral_grid)  # construct model
     simulation = initialize!(model)  
     initialize!(simulation)
@@ -44,11 +43,11 @@
 
     function parameterization_tendencies(diagn, progn, time, model)
         diagn_new = deepcopy(diagn)
-        SpeedyWeather.parameterization_tendencies!(diagn_new, deepcopy(progn), time, deepcopy(model))
+        SpeedyWeather.parameterization_tendencies!(diagn_new, deepcopy(progn), deepcopy(time), deepcopy(model))
         return diagn_new
     end 
 
-    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> parameterization_tendencies(diagn_copy, x, progn.clock.time, model), ddiagn_copy, progn_copy)
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(11,1), x -> parameterization_tendencies(diagn_copy, x, progn.clock.time, model), ddiagn_copy, progn_copy)
     
     # TO-DO this test is broken, they gradients don't line up
     #@test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
@@ -73,6 +72,7 @@
 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> ocean_timestep(progn_copy, x, model), dprogn_copy, diagn_copy)
 
+    # pass
     @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(ddiagn)[1],rtol=1e-4,atol=1e-2))
 
     #
@@ -95,6 +95,7 @@
 
     fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> land_timestep(progn_copy, x, model), dprogn_copy, diagn_copy)
 
+    # pass
     @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(ddiagn)[1],rtol=1e-4,atol=1e-2))
 
     #####
@@ -110,7 +111,7 @@
     ddiag_copy = deepcopy(ddiag)
     progn_copy = deepcopy(progn)
     dprogn = make_zero(progn)
-
+   
     autodiff(Reverse, SpeedyWeather.dynamics_tendencies!, Const, Duplicated(diagn, ddiag), Duplicated(progn, dprogn), Const(lf2), Duplicated(model, make_zero(model)))
 
     function dynamics_tendencies(diagn, progn, lf, model)
@@ -119,15 +120,15 @@
         return diagn_new
     end 
 
-    fd_vjp = FiniteDifferences.j′vp(central_fdm(9,1), x -> dynamics_tendencies(diagn_copy, x, lf2, model), ddiag_copy, progn_copy)
+    fd_vjp = FiniteDifferences.j′vp(central_fdm(5,1), x -> dynamics_tendencies(diagn_copy, x, lf2, model), ddiag_copy, progn_copy)
 
     # there are some NaNs in the FD, that's why this test is currently broken
-    @test all(isapprox.(to_vec(fd_vjp)[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
+    @test all(isapprox.(to_vec(fd_vjp[1])[1], to_vec(dprogn)[1],rtol=1e-4,atol=1e-1))
     
     #
     # Implicit correction 
     #
-
+    # continue here
     diagn_copy = deepcopy(diagn)
     ddiag = make_zero(diagn_copy)
     ddiag_copy = deepcopy(ddiag)
@@ -138,7 +139,7 @@
 
     function implicit_correction(diagn, implicit, progn)
         diagn_new = deepcopy(diagn)
-        SpeedyWeather.implicit_correction(diagn_new, implicit, progn)
+        SpeedyWeather.implicit_correction!(diagn_new, deepcopy(implicit), deepcopy(progn))
         return diagn_new
     end 
 
@@ -150,7 +151,6 @@
     # transform!(diagn, progn, lf2, model)
     #
 
-    fill!(diagn.tendencies, 0, PrimitiveWetModel)
     diag_copy = deepcopy(diagn)
     
     ddiag = one(diagn)
@@ -163,7 +163,7 @@
 
     function transform_diagn(diag, progn, lf2, model)
         diag_copy = deepcopy(diag)
-        transform!(diag_copy, progn, lf2, model)
+        transform!(diag_copy, deepcopy(progn), lf2, deepcopy(model))
         return diag_copy
     end 
 
diff --git a/test/transforms/spectral_transform_ad_rules.jl b/test/transforms/spectral_transform_ad_rules.jl
index 5bc74b621..4afbd6fb5 100644
--- a/test/transforms/spectral_transform_ad_rules.jl
+++ b/test/transforms/spectral_transform_ad_rules.jl
@@ -98,4 +98,8 @@ if VERSION <= v"1.11.0"
         autodiff(Reverse, timestep_oop!, Const, Duplicated(progn_new, dprogn_new), Duplicated(progn, d_progn), Duplicated(diagn, d_diag), Const(dt), Duplicated(model, d_model))
         @test sum(to_vec(d_progn)[1]) != 0
     end 
+else 
+    @testset "Complete Differentiability" begin
+        @test_broken false # we report a broken test here on v1.11, just to indicate that this (properly) doens't work yet
+    end 
 end 
\ No newline at end of file