Compatibility with DifferentialEquations.jl

[JordiBolibar]

Hi, I just came across this package and it looks pretty neat for the sort of work I'm doing. I see that @luraess is adding some parameters and laws related to ice flow modelling, which I'd be interested in using.

My only question here would be: are this sort of structures compatible with DifferentialEquations.jl? The solvers in this library need a function with the following signature: iceflow(du, u, p, t), where p can be a set of parameters to be passed to the solver. Normally a tuple is used, but I've always found this really awkward and hard to track for large sets of parameters. Using GEO_units could be a much more elegant way to deal with this.

Thanks in advance!

[boriskaus]

I don't know whether this is compatible. If you can give us a minimum working example (MWE), we can give it a shot.
The GeoUnit structure is probably not what you want to use (that is simply a structure that has the original value and units @ a separate spot, so we can non-dimensionalize/dimensionalize them).
Yet, the MaterialParams structure is potentially useful, as the main way we use that is in a Tuple. Within the computational routine you would in that case have access to the creep laws & would be able to call the GeoParams computational routines that work on that.

julia> MatParam     = ( SetMaterialParams(Name="Viscous Matrix", Phase=1, 
                                            Density   = ConstantDensity(),
                                            CreepLaws = LinearViscous(η=1e23Pa*s), CharDim=CharDim),
                        SetMaterialParams(Name="Viscous Sinker", Phase=2, 
                                            Density   = PT_Density(),
                                    CompositeRheology = CompositeRheology(LinearViscous(η=1e21Pa*s), DiffusionCreep(), ConstantElasticity()), CharDim=CharDim)
                      )
Phase 1 : Viscous Matrix
        | [non-dimensional units]
        | 
        |-- Density           : Constant density: ρ=2.8999999999999996e-18 
        |-- Gravity           : Gravitational acceleration: g=9.81e20 
        |-- CreepLaws         : Linear viscosity: η=999.9999999999998 
Phase 2 : Viscous Sinker
        | [non-dimensional units]
        | 
        |-- Density           : P/T-dependent density: ρ0=2.8999999999999996e-18, α=0.038194500000000006, β=0.01, T0=0.21454659702313156, P0=0.0 
        |-- Gravity           : Gravitational acceleration: g=9.81e20 
        |-- CompositeRheology : --⟦▪̲̅▫̲̅▫̲̅▫̲̅----⟦▪̲̅▫̲̅▫̲̅▫̲̅----/\/\/--
        
julia> typeof(MatParam)
Tuple{MaterialParams{14, Tuple{ConstantDensity{Float64, Unitful.FreeUnits{(kg, m⁻³·⁰), 𝐌 𝐋⁻³·⁰, nothing}}}, Tuple{ConstantGravity{Float64, Unitful.FreeUnits{(m, s⁻²·⁰), 𝐋 𝐓⁻²·⁰, nothing}}}, Tuple{LinearViscous{Float64, Unitful.FreeUnits{(Pa, s), 𝐌 𝐋⁻¹·⁰ 𝐓⁻¹·⁰, nothing}}}, Tuple{}, Tuple{}, Tuple{}, Tuple{}, Tuple{}, Tuple{}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}, MaterialParams{14, Tuple{PT_Density{Float64, Unitful.FreeUnits{(kg, m⁻³·⁰), 𝐌 𝐋⁻³·⁰, nothing}, Unitful.FreeUnits{(K⁻¹·⁰,), 𝚯⁻¹·⁰, nothing}, Unitful.FreeUnits{(Pa⁻¹·⁰,), 𝐋 𝐓²·⁰ 𝐌⁻¹·⁰, nothing}, Unitful.FreeUnits{(K,), 𝚯, Unitful.Affine{-5463//20}}, Unitful.FreeUnits{(MPa,), 𝐌 𝐋⁻¹·⁰ 𝐓⁻²·⁰, nothing}}}, Tuple{ConstantGravity{Float64, Unitful.FreeUnits{(m, s⁻²·⁰), 𝐋 𝐓⁻²·⁰, nothing}}}, Tuple{}, Tuple{}, Tuple{}, Tuple{CompositeRheology{Tuple{LinearViscous{Float64, Unitful.FreeUnits{(Pa, s), 𝐌 𝐋⁻¹·⁰ 𝐓⁻¹·⁰, nothing}}, DiffusionCreep{Float64, 0, Unitful.FreeUnits{, NoDims, nothing}, Unitful.FreeUnits{(m³·⁰, MPa⁻¹·⁰, s⁻¹·⁰), 𝐋⁴·⁰ 𝐓 𝐌⁻¹·⁰, nothing}, Unitful.FreeUnits{(kJ, mol⁻¹·⁰), 𝐋²·⁰ 𝐌 𝐍⁻¹·⁰ 𝐓⁻²·⁰, nothing}, Unitful.FreeUnits{(m³·⁰, mol⁻¹·⁰), 𝐋³·⁰ 𝐍⁻¹·⁰, nothing}, Unitful.FreeUnits{(J, K⁻¹·⁰, mol⁻¹·⁰), 𝐋²·⁰ 𝐌 𝐍⁻¹·⁰ 𝚯⁻¹·⁰ 𝐓⁻²·⁰, nothing}}, ConstantElasticity{Float64, Unitful.FreeUnits{(Pa,), 𝐌 𝐋⁻¹·⁰ 𝐓⁻²·⁰, nothing}, Unitful.FreeUnits{, NoDims, nothing}}}, 3, 0, Bool[0, 0, 0], 0, Bool[0, 0, 0], 0, Bool[0, 0, 0]}}, Tuple{}, Tuple{}, Tuple{}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}}

[albert-de-montserrat]

Thanks a lot for the interest @JordiBolibar ! I'm not very familiar with how DiffEq works, but I doubt that it works with GeoParams out-of-the-box. It would be very intersting to have them work together though, if you can provide a minimal example of what you would like to do, I can try to make an interface.

In case this is of interest for you, GeoParams is currently compatible with ForwardDiff.jl and (with some care) ParallelStencil.jl.

[JordiBolibar]

Hi @boriskaus and @albert-de-montserrat, thanks for the feedback.

Indeed, I still need to properly understand the intended goals of your library and how good a match would be with my problem. My main interests in order to use this for now are the following:

• Have a tidy data structure to gather all the parameters with units used in the forward ice flow model.
• Hopefully be able to use some of your creep laws for ice, which could be computationally optimized.

Most likely, two main technical constraints lie ahead:

• Compatibility with DifferentialEquations.jl
• Compatibility with AD frameworks, namely Zygote.jl and Enzyme.jl. For now we're working with Zygote, but we plan to soon move to Enzyme once it becomes more stable.

As requested, here's a MWE that might be useful to illustrate this. What would be interesting to update here is the context tuple that is passed to the solver. Many of the parameters (e.g. g, or rho) are declared as global constants. Ideally, if I had a better data structure than a tuple it would make much more sense to thread everything locally to avoid inconsistencies.

Not sure how feasible is all this, but I'd love to be able to re-use this sort of infrastructure.

[Iddingsite]

Hi everyone, I will add a bit on this (already a bit old) discussion as I have some experience using DifferentialEquations and a bit using GeoParams too.
I can tell you that there is no issue combining these two packages together. To prove my point, I wrote a small code solving the 2D non-linear heat diffusion equation using a solver from DiffEq and updating the heat capacity and thermal conductivity using GeoParams routines. I am also using ParallelStencil to parallelize it and make it GPU compatible but you can put that aside.

using GeoParams
using DifferentialEquations
using ParallelStencil
using Plots
using BenchmarkTools
using Logging: global_logger
using TerminalLoggers: TerminalLogger
global_logger(TerminalLogger())

println("Thread number: $(Threads.nthreads())")

# initialise ParallelStencil
const USE_GPU = false
@static if USE_GPU
    @init_parallel_stencil(CUDA, Float64, 2);
else
    @init_parallel_stencil(Threads, Float64, 2);
end

# take the average of the diffusion coefficient btw 2 cells in each direction
macro av_D_x(D, ix, iz) esc(:((0.5*($D[$iz,$ix]+$D[$iz,$ix+1])))) end
macro av_D_z(D, ix, iz) esc(:((0.5*($D[$iz,$ix]+$D[$iz+1,$ix])))) end

# macros for flux discretisation
macro qx(D,T,ix,iz,Δx_) esc(:((@av_D_x($D,$ix,$iz)) * ($T[$iz,$ix+1]-$T[$iz,$ix])*$Δx_ )) end
macro qz(D,T,ix,iz,Δz_) esc(:((@av_D_z($D,$ix,$iz)) * ($T[$iz+1,$ix]-$T[$iz,$ix])*$Δz_ )) end

# discretisation in space
@parallel_indices (iz, ix) function stencil_diffusion!(dtT, T, k, Cp, ρ, Δx_, Δz_)
    # iterate inside the arrays and discretised the diffusion term
    if ix>1 && ix<size(dtT,2) && iz>1 && iz<size(dtT,1)
        dtT[iz,ix] = 1/(ρ[iz,ix]*Cp[iz,ix]) *
                     ((@qx(k,T,ix,iz,Δx_) - @qx(k,T,ix-1,iz,Δx_)) * Δx_ +
                      (@qz(k,T,ix,iz,Δz_) - @qz(k,T,ix,iz-1,Δz_)) * Δz_)
    end
    return nothing
end

# discretisation in space using ParallelStencil and update of heatcapacity
function semi_discretisation(dTdt,T,p,t)

    geoparam, thermal_param = p
    MatParam, phase = geoparam
    k, Cp, ρ = thermal_param

    # # update Cp
    compute_heatcapacity!(Cp, MatParam, phase, (;T=T))

    @parallel stencil_diffusion!(dTdt, T, k, Cp, ρ, Δx_, Δz_)
end

# grid
const Lz = 20_000 # in m
const Lx = 20_000 # in m
const nx = 400
const nz = 400
const Δx = Lx / (nx)
const Δx_ = 1 / Δx  # to avoid division in the discretisation
const Δz = Lz / (nz)
const Δz_ = 1 / Δz  # to avoid division in the discretisation
const x = range(Δx/2, length=nx, stop= Lx-Δx/2)
const z = range(Δz/2, length=nx, stop= Lz-Δz/2)
const grid = (x=x' .* @ones(nz), z=(@ones(nx))' .* z)

# total time in Myrs
const tfinal = 0.5 * 1e6 * 3600 * 24 * 365.25  # 0.5 Myrs in total in s

# initial conditions
T0 = @zeros(nz, nx)

const x0 = Lx/2  # x coordinate of the center of the cylinder
const z0 = Lz/2  # z coordinate of the center of the cylinder
const R = 50*Δx   # radius of the cylinder

# define centered cylinder anomaly as initial conditions in K
for I in CartesianIndices(T0)
    r = sqrt((grid[1][I] - x0)^2 + (grid[2][I] - z0)^2)
    if r <= R
        T0[I] = 1073.15
    else
        T0[I] = 873.15
    end
end

# define material parameters using GeoParams
MatParam = (SetMaterialParams(;Name="Crust", Phase=1,
                        Density   = ConstantDensity(ρ=2700kg/m^3),
                        HeatCapacity = T_HeatCapacity_Whittington(),
                        Conductivity = T_Conductivity_Whittington()),)

# define phase
phase = ones(Int64,nz,nx)
# define density
ρ=similar(T0)
compute_density!(ρ, MatParam, phase, (; T=T0))
# define specific heat capacity
Cp = similar(T0)
compute_heatcapacity!(Cp, MatParam, phase, (; T=T0))
# define thermal conductivity
k = similar(T0)
compute_conductivity!(k, MatParam, phase, (; T=T0))


# callback function to update conductivity using GeoParams (can't update it directly in discretisation function because 1/T in thermal diffusivity definition)
function thermal_conductivity_call_func(T, t, integrator)

    geoparam, thermal_param = integrator.p
    MatParam, phase = geoparam
    k, Cp, ρ = thermal_param

    compute_conductivity!(k, MatParam, phase, (;T=T))
end

# create a callback function to call at each time step for the thermal conductivity
thermal_param_call = FunctionCallingCallback(thermal_conductivity_call_func; funcat=Vector{Float64}(), func_everystep=true, func_start = false, tdir=1);

# define a named tuple containing the parameters for the ODE solver
p = (geoparam=(MatParam, phase), thermal_param=(k, Cp, ρ))
# define the initial and final time
t = [0, tfinal]

# define the ODE problem
prob = ODEProblem(semi_discretisation, T0, t, p);

# solve with ROCK2 (stabilized explicit solver) using DiffEq


@time sol = solve(prob, ROCK2(), progress = true, progress_steps = 1, save_start=true,callback= thermal_param_call, abstol=1e-6,reltol=1e-6);
# before compilation, with 8 threads on my machine
# 19.775351 seconds (21.22 M allocations: 2.868 GiB, 5.18% gc time, 49.08% compilation time)

@time sol = solve(prob, ROCK2(), progress = true, progress_steps = 1, save_start=true, callback= thermal_param_call, abstol=1e-6,reltol=1e-6);
# after compilation, with 8 threads on my machine
# 9.916377 seconds (432.39 k allocations: 1.636 GiB, 3.11% gc time)


l = @layout [a ; b]
p1 = heatmap(x.*1e-3,z.*1e-3,sol[1][:,:].-273.15,aspect_ratio=:equal, c=:heat, xlim=(Δx/2*1e-3,(Lx-Δx/2)*1e-3),zlim=(Δz/2*1e-3,(Lz-Δz/2)*1e-3), clim=(minimum(T0.-273.15), maximum(T0.-273.15)),xticks = 0:5:Lx, title="t = 0", colorbar_title="Temperature (°C)",ylabel = "Distance (km)")
p2 = heatmap(x.*1e-3,z.*1e-3,sol[end][:,:].-273.15,aspect_ratio=:equal, c=:heat, xlim=(Δx/2*1e-3,(Lx-Δx/2)*1e-3),zlim=(Δz/2*1e-3,(Lz-Δz/2)*1e-3), clim=(minimum(T0.-273.15), maximum(T0.-273.15)),xticks = 0:5:Lx, title="t = $(tfinal/(1e6 * 3600 * 24 * 365.25)) Ma", colorbar_title="Temperature (°C)", xlabel = "Distance (km)", ylabel = "Distance (km)")
display(plot(p1, p2, layout = l,aspect_ratio=:equal))

This take 10 sec to solve on 8 threads on my computer for 400x400 nodes the cooling of a cylindrical thermal anomaly during 0.5 Ma which is quite cool. I also didn't parallelized the GeoParam routines so there is still room for another gain in speed. It outputs this:

Note that I had to create a callback function to update the thermal conductivity because of a 1/T in its definition, which is not compatible with the solver I am using (a stabilized explicit method). Normally you could put everything together in the semidiscretisation function.

Concerning Zygote and Enzyme, I don't have experience using them so I don't know. Would be worth a try maybe? Anyway, there is still no Ice creep law implemented yet, so I guess you will have to wait a bit more on that side to try anything 😄

[albert-de-montserrat]

That's great @Iddingsite thanks for sharing!

Quick update: it's still WIP but most of the GeoParams.jl kernels are already compatible with the AD packages.

On a side note, instead of using macros inside ParallelStencil.jl kernels, I find that it's usually better to use closures:

@parallel_indices (iz, ix) function stencil_diffusion!(dtT, T, D, Cp, ρ, Δx_, Δz_)
    # take the average of the diffusion coefficient btw 2 cells in each direction
    Base.@propagate_inbounds @inline av_D_x(ix, iz) = 0.5*(D[iz,ix]+D[iz,ix+1])
    Base.@propagate_inbounds @inline av_D_z(ix, iz) = 0.5*(D[iz,ix]+D[iz+1,ix])
    # macros for flux discretisation
    Base.@propagate_inbounds @inline qx(ix,iz) = av_D_x(ix,iz) * (T[iz,ix+1]-T[iz,ix])*Δx_ 
    Base.@propagate_inbounds @inline qz(ix,iz) = av_D_z(ix,iz) * (T[iz+1,ix]-T[iz,ix])*Δz_ 

    # iterate inside the arrays and discretised the diffusion term
    @inbounds if  1<ix<size(dtT,2) && 1<iz<size(dtT,1)
        dtT[iz,ix] = 1/(ρ[iz,ix]*Cp[iz,ix]) *
                     ((qx(ix,iz) - qx(ix-1,iz)) * Δx_ +
                      (qz(ix,iz) - qz(ix,iz-1)) * Δz_)
    end
    return nothing
end

This yields about 15% extra speed in my machine.

[luraess]

This yields about 15% extra speed in my machine.

Do you observe this on CPU or GPU (or both)?

[Iddingsite]

Interesting @albert-de-montserrat, I didn't know this notation. I think it makes the code cleaner that way actually!

But I couldn't reproduce your performance. With closures, it is about 10% slower on my computer 😅 I also have slightly more allocations. I checked with @btime. I guess it is quite heavy dependent on the processor if you were talking about CPU.

Also, I forgot to add the save_everystep argument in the solve function in my previous code. It prevents saving every timestep so it should reduce the allocations.

# after compilation, with 8 threads on my machine
@btime sol = solve(prob, ROCK2(), progress = true, progress_steps = 1, save_start=true, save_everystep=false, callback= thermal_param_call, abstol=1e-6,reltol=1e-6);
# with macros
# 9.194 s (459740 allocations: 60.82 MiB)
# with closures
# 9.457 s (462624 allocations: 60.90 MiB)

EDIT

After trying in the terminal and not in Vscode, it is slightly faster with closures (not sure if the difference is real here, I guess it would be worth to try with a longer model and a higher resolution but it is time to sleep for me):

# with macros
# 9.144 s (476026 allocations: 60.68 MiB)
# with closures
# 9.112 s (476449 allocations: 60.69 MiB)

So I don't know anymore 😁

[JordiBolibar]

This is awesome @Iddingsite! This is one more step towards my ideal setup. Right now I'm using:

• DifferentialEquations.jl solvers
• Zygote.jl and ReverseDiff.jl for AD
• Really basic data structures to ensure that AD works in my problem

Which I would like to improve towards:

• DifferentialEquations.jl solvers
• Enzyme.jl or Zygote.jl for AD
• GeoParams.jl for passing parameters to the PDE
• ParallelStencil.jl or ImplicitGlobalGrid.jl to parallelize the discretized ODEs

I would say the tricky thing here is the AD. Without AD you just proved that everything works smoothly, so I guess it's getting closer! 😄

[boriskaus]

If you are interested in combining Enzyme.jl with ParallelStencil.jl/ImplicitGlobalGrid.jl, have a look at these examples.

Using AD with GeoParams.jl is certainly on the todo list; we simply haven't gotten around doing it (getting funding to work on this would help; proposal was submitted yesterday...).

[albert-de-montserrat]

This yields about 15% extra speed in my machine.

Do you observe this on CPU or GPU (or both)?

Haven't done a thorough benchmark but in both. In some cases there's barely any difference though, but makes the code a bit easier to read in my opinion.

[albert-de-montserrat]

Interesting @albert-de-montserrat, I didn't know this notation. I think it makes the code cleaner that way actually!

But I couldn't reproduce your performance. With closures, it is about 10% slower on my computer 😅 I also have slightly more allocations. I checked with @btime. I guess it is quite heavy dependent on the processor if you were talking about CPU.

Also, I forgot to add the save_everystep argument in the solve function in my previous code. It prevents saving every timestep so it should reduce the allocations.

# after compilation, with 8 threads on my machine
@btime sol = solve(prob, ROCK2(), progress = true, progress_steps = 1, save_start=true, save_everystep=false, callback= thermal_param_call, abstol=1e-6,reltol=1e-6);
# with macros
# 9.194 s (459740 allocations: 60.82 MiB)
# with closures
# 9.457 s (462624 allocations: 60.90 MiB)

EDIT

After trying in the terminal and not in Vscode, it is slightly faster with closures (not sure if the difference is real here, I guess it would be worth to try with a longer model and a higher resolution but it is time to sleep for me):

# with macros
# 9.144 s (476026 allocations: 60.68 MiB)
# with closures
# 9.112 s (476449 allocations: 60.69 MiB)

So I don't know anymore 😁

At least is not slower :) I was timing it on my laptop with @time yesterday so maybe it wasn't very accurate. This is what I see now on my worksation (I made the problem smaller, 32x32), so effectively no difference:

julia> @btime solve($prob, $ROCK2(), progress = true, progress_steps = 1, save_start=true,callback= thermal_param_call, abstol=1e-6,reltol=1e-6);
  1.384 ms (7588 allocations: 1.56 MiB)
julia> @btime solve($prob2, $ROCK2(), progress = true, progress_steps = 1, save_start=true,callback= thermal_param_call, abstol=1e-6,reltol=1e-6);
  1.306 ms (7591 allocations: 1.56 MiB)

where prob2 is the version with closures. Could you benchmark it on a GPU once you have everything running with ParallelStencil?

EDIT: had an error in the benchmark, version with closures is consistently faster, but just around 5%

[JordiBolibar]

If you are interested in combining Enzyme.jl with ParallelStencil.jl/ImplicitGlobalGrid.jl, have a look at these examples.

Using AD with GeoParams.jl is certainly on the todo list; we simply haven't gotten around doing it (getting funding to work on this would help; proposal was submitted yesterday...).

Thanks @boriskaus. Last week I was at ETH and I had a lengthy chat about this with @utkinis. This is definitely on the to-do list of my project. Hopefully I'll try to join one of the hackweeks in the Black Forest to implement this in the future!

[Iddingsite]

At least is not slower :) I was timing it on my laptop with @time yesterday so maybe it wasn't very accurate. This is what I see now on my worksation (I made the problem smaller, 32x32), so effectively no difference:

julia> @btime solve($prob, $ROCK2(), progress = true, progress_steps = 1, save_start=true,callback= thermal_param_call, abstol=1e-6,reltol=1e-6);
  1.384 ms (7588 allocations: 1.56 MiB)
julia> @btime solve($prob2, $ROCK2(), progress = true, progress_steps = 1, save_start=true,callback= thermal_param_call, abstol=1e-6,reltol=1e-6);
  1.306 ms (7591 allocations: 1.56 MiB)

where prob2 is the version with closures. Could you benchmark it on a GPU once you have everything running with ParallelStencil?

EDIT: had an error in the benchmark, version with closures is consistently faster, but just around 5%

Cool! Then I will start adopting this syntax and I agree that it is also easier to read. I will try to benchmark on GPU next week 🙌