Sensitivity Analysis Failing

[maximilian-gelbrecht]

To demonstrate basic differentiability we aim to do a sensitivity analysis, e.g. by computing the derivative $$\frac{\partial T_i(N)}{\partial \mathbf{X}(0)}$$ where $$T_i$$ is the temperature at some grid point $$i$$ evolved by $$N$$ days from initial condition $$\mathbf{X}(0)$$

Currently a simplified example of this is failing. I did try several somewhat different version of writing the vorticity_at_gridpoint! and all of them failed (some with a different error). As @bgroenks96 pointed out, we could already now do exactly this kind of sensitivity analysis by using a one-hot-array as the seed when directly differentiating timestep! (which works). However, when an error already occurs here, it'll also cause other problems in other things we want to look at.

using SpeedyWeather, Enzyme 

# first we initialize a model 
# for now, we pick our least complex model, to make everything run faster
arch = SpeedyWeather.CPU()
spectral_grid = SpectralGrid(trunc=8, nlayers=1, architecture=arch)
model = BarotropicModel(spectral_grid)
simulation = initialize!(model)
initialize!(simulation)
run!(simulation, period=Day(3)) # a bit of spin-up

const i_point = 12 # pick this kind of random point 
const dt = model.time_stepping.Δt

# in this example, to demonstrate the error, we just take some point of the vorticity field
# that's actually spectral field, which doesn't make too much sense for a sensitivity analysis, 
# but this is just to demonstrate the error with Enzyme
function vorticity_at_gridpoint!(progn, diagn, model)
    # just do a single timestep here for this example, later this will be replaced with run!(... , period=Day(N_days))
    SpeedyWeather.timestep!(progn, diagn, dt, model)
    return abs(progn.vor[i_point, 1, end]) 
end 

progn = simulation.prognostic_variables
dprogn = zero(progn)

diagn = simulation.diagnostic_variables
ddiag = make_zero(diagn)

# just to trigger some compilation
vorticity_at_gridpoint!(progn, diagn, model)

autodiff(Reverse, vorticity_at_gridpoint!, Active, Duplicated(progn, dprogn), Duplicated(diagn, ddiag), Const(model))

# catch the error to display it
try
    autodiff(Reverse, vorticity_at_gridpoint!, Active, Duplicated(progn, dprogn), Duplicated(diagn, ddiag), Const(model))
catch err
    println("Enzyme error caught")
end

# contrary, this works fine: 
autodiff(Reverse, SpeedyWeather.timestep!, Const, Duplicated(progn, dprogn), Duplicated(diagn, ddiag), Const(dt), Const(model))

Yields currently a method error. Last time I tried this two months ago, it was something else. Now it looks like a rule is missing for something? It's the activity analysis determining Const activity for the output of our fourier transform, where our rule is just for Duplicated.

ERROR: MethodError: no method matching augmented_primal(::EnzymeCore.EnzymeRules.RevConfigWidth{…}, ::Const{…}, ::Type{…}, ::Const{…}, ::Const{…}, ::Duplicated{…}, ::Const{…})

Closest candidates are:
  augmented_primal(::EnzymeCore.EnzymeRules.RevConfig, ::Const{typeof(Enzyme.pmap)}, ::Type{Const{Nothing}}, ::BodyTy, ::Any, ::Annotation...) where {BodyTy, N}
   @ Enzyme ~/.julia/packages/Enzyme/12QGc/src/internal_rules.jl:337
  augmented_primal(::EnzymeCore.EnzymeRules.RevConfig, ::Const{<:KernelAbstractions.Kernel}, ::Type{Const{Nothing}}, ::Any...; ndrange, workgroupsize) where N
   @ EnzymeExt ~/.julia/packages/KernelAbstractions/lGrz7/ext/EnzymeCore08Ext.jl:214
  augmented_primal(::Any, ::Const{typeof(QuadGK.quadgk)}, ::Type{RT}, ::Any, ::Annotation{T}...; kws...) where {RT, T}
   @ QuadGKEnzymeExt ~/.julia/packages/QuadGK/7rND3/ext/QuadGKEnzymeExt.jl:6
  ...

Some type information was truncated. Use `show(err)` to see complete types.

The code is on the current main with Julia v1.10 (there are currently still problems with Julia v1.11).

@michel2323

[maximilian-gelbrecht]

I just quickly looked at this myself now, and the issue is related to our custom rule for _fourier! that's missing with different activities.

MethodError: no method matching augmented_primal(::EnzymeCore.EnzymeRules.RevConfigWidth{1, false, false, (false, true, true, true, true), false, false}, ::Const{typeof(SpeedyWeather.SpeedyTransforms._fourier!)}, ::Type{Const{Nothing}}, ::Const{Array{ComplexF32, 3}}, ::Const{Array{ComplexF32, 3}}, ::Duplicated{Field{Float32, 2, Matrix{Float32}, OctahedralGaussianGrid{SpeedyWeather.Architectures.CPU{KernelAbstractions.CPU}, Vector{UnitRange{Int64}}, Vector{Int64}}}}, ::Const{SpectralTransform{Float32, SpeedyWeather.Architectures.CPU{KernelAbstractions.CPU}, Array, Spectrum{SpeedyWeather.Architectures.CPU{KernelAbstractions.CPU}, Vector{UnitRange{Int64}}, Vector{Int64}}, OctahedralGaussianGrid{SpeedyWeather.Architectures.CPU{KernelAbstractions.CPU}, Vector{UnitRange{Int64}}, Vector{Int64}}, Vector{Float32}, Matrix{Int64}, Vector{ComplexF32}, Matrix{ComplexF32}, Array{ComplexF32, 3}, LowerTriangularMatrix{Float32, Vector{Float32}, Spectrum{SpeedyWeather.Architectures.CPU{KernelAbstractions.CPU}, Vector{UnitRange{Int64}}, Vector{Int64}}}, LowerTriangularArray{Float32, 2, Matrix{Float32}, Spectrum{SpeedyWeather.Architectures.CPU{KernelAbstractions.CPU}, Vector{UnitRange{Int64}}, Vector{Int64}}}}})

So here the scratch memory that holds the fourier transformed but not yet legendre transformed arrays is marked with Const in this case instead of Duplicated activity. I am a bit confused about this, because I think it should be Duplicated as far as I can see it, and it is like that if I only differentiate timestep! by itself.

[maximilian-gelbrecht]

Thanks @vchuravy already for helping out!

When marking model Duplicated the example above works. So, there's still some issue with scratch memory.

So I think either, it's that we are not really using the scratch memory that's already implemented in DiagnosticVariables correctly, or maybe it's related to what's addressed in #791 (that there's still some scratch memory actually remaining in the SpectralTransform)

[maximilian-gelbrecht]

One problem fixed in #828 , but this still needs a little more work to actually be a working sensitivity analysis. I already started adjusting the script a little. I'll keep this issue open in case other problems occur during the process.

[maximilian-gelbrecht]

So, my adjustment currently causes a StackOverflowError. I didn't investigate this yet at all, but there still some issues to resolve here

using SpeedyWeather, Enzyme, Makie

# first we initialize a model 
# for now, we pick our least complex model, to make everything run faster
arch = SpeedyWeather.CPU()
spectral_grid = SpectralGrid(trunc=32, nlayers=8, architecture=arch)
model = PrimitiveWetModel(spectral_grid)
simulation = initialize!(model)
initialize!(simulation)
run!(simulation, period=Day(5)) # a bit of spin-up

const i_point = 443 # pick this kind of random point (it's ≈ Berlin)
const dt = model.time_stepping.Δt

# in this example, to demonstrate the error, we just take some point of the vorticity field
# that's actually spectral field, which doesn't make too much sense for a sensitivity analysis, 
# but this is just to demonstrate the error with Enzyme
function temperature_at_gridpoint!(progn, diagn, model)
    # just do a single timestep here for this example, later this will be replaced with run!(... , period=Day(N_days))
    SpeedyWeather.timestep!(progn, diagn, dt, model)
    return diagn.grid.temp_grid[i_point, 8]
end 

progn = simulation.prognostic_variables
dprogn = zero(progn)

diagn = simulation.diagnostic_variables
ddiag = make_zero(diagn)

# just to trigger some compilation
temperature_at_gridpoint!(progn, diagn, model)

autodiff(Reverse, temperature_at_gridpoint!, Active, Duplicated(progn, dprogn), Duplicated(diagn, ddiag), Const(model))