Computing parameter sensitivities with autodiff

This example demonstrates how to parameterize and differentiate a simulation with two parameters (summer and winter n-factors) using forward-mode automatic simulation.

TODO: add more detail/background

using CryoGrid

Set up forcings and boundary conditions similarly to other examples:

forcings = loadforcings(CryoGrid.Forcings.Samoylov_ERA_obs_fitted_1979_2014_spinup_extended_2044);
soilprofile, tempprofile = CryoGrid.SamoylovDefault
freezecurve = PainterKarra(swrc=VanGenuchten())
soilprofile = SoilProfile(0.0u"m" => SimpleSoil(; freezecurve))
grid = CryoGrid.DefaultGrid_5cm
initT = initializer(:T, tempprofile)
tile = CryoGrid.SoilHeatTile(
    :T,
    TemperatureBC(Input(:Tair), NFactor(nf=Param(0.5), nt=Param(0.9))),
    GeothermalHeatFlux(0.053u"W/m^2"),
    soilprofile,
    forcings,
    initT;
    grid=grid
)
tspan = (DateTime(2010,9,1),DateTime(2011,10,1))
u0, du0 = @time initialcondition!(tile, tspan);

  8.613853 seconds (5.81 M allocations: 408.696 MiB, 0.93% gc time, 100.00% compilation time)

We can retrieve the parameters of the system from tile:

para = CryoGrid.parameters(tile)

CryoGridParams{Float64} with 10 parameters
Parameters:
┌────────────┬───────────┬─────────┬────────┬───────────────────────────────────
│  component │ fieldname │     val │  layer │                                  ⋯
│   UnionAll │    Symbol │ Float64 │ Symbol │ U{Nothing, ClosedInterval{Float6 ⋯
├────────────┼───────────┼─────────┼────────┼───────────────────────────────────
│    NFactor │        nf │     0.5 │    top │                                  ⋯
│    NFactor │        nt │     0.9 │    top │                                  ⋯
│ SimpleSoil │       por │     0.5 │ ground │                                  ⋯
│ SimpleSoil │       sat │     1.0 │ ground │                                  ⋯
│ SimpleSoil │       org │     0.0 │ ground │                                  ⋯
│ NamedTuple │      kh_m │     3.8 │ ground │                                  ⋯
│ NamedTuple │      kh_o │    0.25 │ ground │                                  ⋯
│ NamedTuple │      ch_m │   2.0e6 │ ground │                                  ⋯
│ NamedTuple │      ch_o │   2.5e6 │ ground │                                  ⋯
│ NamedTuple │    kw_sat │  1.0e-5 │ ground │                                  ⋯
└────────────┴───────────┴─────────┴────────┴───────────────────────────────────
                                                               5 columns omitted

Create the CryoGridProblem.

prob = CryoGridProblem(tile, u0, tspan, saveat=3600.0)

CryoGridProblem with uType ComponentArrays.ComponentVector{Float64, Vector{Float64}, Tuple{ComponentArrays.Axis{(T = 1:1:218,)}}} and tType Float64. In-place: true
Non-trivial mass matrix: false
timespan: (6.34505184e10, 6.34846464e10)
u0: ComponentVector{Float64}(T = [-1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0  …  -9.0, -7.933333333333334, -5.800000000000001, -3.6666666666666665, -1.5333333333333337, 0.5999999999999996, 2.733333333333334, 4.866666666666666, 7.0, 9.133333333333333])

Solve the forward problem with default parameter settings:

sol = @time solve(prob)
out = CryoGridOutput(sol)

CryoGridOutput with 9481 time steps (2010-09-01T00:00:00 to 2011-10-01T00:00:00) and 1 variable:
    T => DimArray of Quantity{Float64, 𝚯, Unitful.FreeUnits{(K,), 𝚯, Unitful.Affine{-5463//20}}} with dimensions (218, 9481)

Import relevant packages for automatic differentiation.

using ForwardDiff
using SciMLSensitivity
using Zygote

Define a "loss" function; here we'll just take the mean over the final temperature field.

using OrdinaryDiffEq
using Statistics
function loss(prob::CryoGridProblem, p)
    newprob = remake(prob, p=p)
    # Here we specify the sensitivity algorithm. Note that this is only
    # necessary for reverse-mode autodiff with Zygote.
    # autojacvec = true uses ForwardDiff to calculate the jacobian;
    # enabling checkpointing (theoretically) reduces the memory cost of the backwards pass.
    sensealg = InterpolatingAdjoint(autojacvec=true, checkpointing=true)
    newsol = solve(newprob, Euler(), dt=300.0, sensealg=sensealg);
    newout = CryoGridOutput(newsol)
    return mean(ustrip.(newout.T[:,end]))
end

loss (generic function with 1 method)

Compute gradient with forward-mode autodiff:

pvec = vec(prob.p)
fd_grad = @time ForwardDiff.gradient(pᵢ -> loss(prob, pᵢ), pvec)

# We can also try with reverse-mode autodiff. This is generally slower for smaller numbers
# of parmaeters (<100) but could be worthwhile for model configurations with high-dimensional
# parameterizations.
# zy_grad = @time Zygote.gradient(pᵢ -> loss(prob, pᵢ), pvec)
# @assert maximum(abs.(fd_grad .- zy_grad)) .< 1e-6 "Forward and reverse gradients don't match!"

This page was generated using Literate.jl.