Functional inversion tutorial

This tutorial provides a simple example on how to perform a functional inversion using Universal Differential Equations (UDEs) in ODINN.jl. For this, we will generate a synthetic dataset using a forward simulation, and then we will use this dataset to perform the functional inversion. The goal of this functional inversion will be to learn a synthetic law that maps A, i.e. the ice rigidity, to long-term changes in atmospheric surface temperature.

For more details on the functional inversion concept, please refer to the Functional Inversion section in the Optimization page.

Running the whole code

using ODINN

# Define the working directory
working_dir = joinpath(ODINN.root_dir, "demos")

## We fetch the paths with the files for the available glaciers on disk
rgi_paths = get_rgi_paths()

# Ensure the working directory exists
mkpath(working_dir)

# Define which glacier RGI IDs we want to work with
rgi_ids = ["RGI60-11.03638", "RGI60-11.01450", "RGI60-11.02346", "RGI60-07.00065"]
# Define the time step for the simulation output and for the adjoint calculation. In this case, a month.
δt = 1/12

params = Parameters(
    simulation = SimulationParameters(
        working_dir=working_dir,
        use_MB=false,
        use_velocities=true,
        tspan=(2010.0, 2015.0),
        step=δt,
        multiprocessing=true,
        workers=4,
        test_mode=false,
        rgi_paths=rgi_paths,
        gridScalingFactor=4), # Downscale the glacier grid to speed-up this example
    hyper = Hyperparameters(
        batch_size=length(rgi_ids), # We set batch size equals all datasize so we test gradient
        epochs=[15,10],
        optimizer=[ODINN.ADAM(0.01), ODINN.LBFGS(linesearch = ODINN.LineSearches.BackTracking(iterations = 5))]),
    physical = PhysicalParameters(
        minA = 8e-21,
        maxA = 8e-17),
    UDE = UDEparameters(
        optim_autoAD=ODINN.NoAD(),
        grad=ContinuousAdjoint(),
        optimization_method="AD+AD",
        empirical_loss_function = LossH() # Loss function based on ice thickness
    ),
    solver = Huginn.SolverParameters(
        step=δt,
        save_everystep=true,
        progress=true)
)

# We define a synthetic law to generate the synthetic dataset. For this, we use some tabular data from Cuffey and Paterson (2010).
A_law = CuffeyPaterson()

model = Huginn.Model(
    iceflow = SIA2Dmodel(params; A=A_law),
    mass_balance = TImodel1(params; DDF=6.0/1000.0, acc_factor=1.2/1000.0),
)

# We initialize the glaciers with all the necessary data
glaciers = initialize_glaciers(rgi_ids, params)

# Time snapshots for transient inversion
tstops = collect(2010:δt:2015)

# We generate the synthetic dataset using the forward simulation. This will generate a dataset with the ice thickness and surface velocities for each glacier at each time step. The dataset will be used to train the machine learning model.
glaciers = generate_ground_truth(glaciers, params, model, tstops)

## After this forward simulation, we restart the iceflow model to be ready for the inversions
nn_model = NeuralNetwork(params)
A_law = LawA(nn_model, params)
model = Model(
    iceflow = SIA2Dmodel(params; A=A_law),
    mass_balance = TImodel1(params; DDF=6.0/1000.0, acc_factor=1.2/1000.0),
    regressors = (; A=nn_model)
)

# We specify the type of simulation we want to perform
functional_inversion = FunctionalInversion(model, glaciers, params)

# And finally, we just run the simulation
run!(functional_inversion)
[ Info: Filtering out these glaciers from RGI ID list: Any[]
      From worker 4:	Getting raw climate data for: RGI60-07.00065
[ Info: Running forward in-place PDE ice flow model
      From worker 4:	Processing glacier RGI60-11.03638 for PDE forward simulation
      From worker 2:	Processing glacier RGI60-11.01450 for PDE forward simulation
      From worker 3:	Processing glacier RGI60-11.02346 for PDE forward simulation
      From worker 4:	Processing glacier RGI60-07.00065 for PDE forward simulation

Progress:  50%|████████████████████▌                    |  ETA: 0:00:02
Progress:  75%|██████████████████████████████▊          |  ETA: 0:00:01
Progress: 100%|█████████████████████████████████████████| Time: 0:00:02
Running training of UDE...

[ Info: Training with ADAM optimizer
[ Info: Training with custom ContinuousAdjoint{Float64, Int64, DiscreteVJP} method
Iteration: [    1 /    25]     Loss:7.94252e+01
Iteration: [    2 /    25]     Loss:6.57928e+01     Improvement: -17.16 %
Iteration: [    3 /    25]     Loss:5.49244e+01     Improvement: -16.52 %
Iteration: [    4 /    25]     Loss:4.63710e+01     Improvement: -15.57 %
Iteration: [    5 /    25]     Loss:3.96132e+01     Improvement: -14.57 %
Iteration: [    6 /    25]     Loss:3.41950e+01     Improvement: -13.68 %
Iteration: [    7 /    25]     Loss:2.97769e+01     Improvement: -12.92 %
Iteration: [    8 /    25]     Loss:2.61246e+01     Improvement: -12.27 %
Iteration: [    9 /    25]     Loss:2.30760e+01     Improvement: -11.67 %
Iteration: [   10 /    25]     Loss:2.05147e+01     Improvement: -11.10 %
Iteration: [   11 /    25]     Loss:1.83530e+01     Improvement: -10.54 %
Iteration: [   12 /    25]     Loss:1.65219e+01     Improvement: -9.98 %
Iteration: [   13 /    25]     Loss:1.49655e+01     Improvement: -9.42 %
Iteration: [   14 /    25]     Loss:1.36379e+01     Improvement: -8.87 %
Iteration: [   15 /    25]     Loss:1.25008e+01     Improvement: -8.34 %
Iteration: [   16 /    25]     Loss:1.25008e+01     Improvement: 0.00 %
[ Info: Training with BFGS optimizer
[ Info: Training with custom ContinuousAdjoint{Float64, Int64, DiscreteVJP} method
Iteration: [   17 /    25]     Loss:1.25008e+01     Improvement: 0.00 %
Iteration: [   18 /    25]     Loss:1.04657e+01     Improvement: -16.28 %

Step-by-step explanation of the tutorial

Here we will cover in detail each one of the steps that lead us to run the FunctionalInversion from the previous example. The goal of this simple example is to learn a mapping of a law for A, the creep coefficient of ice. Mathematically, we make A depends on the long term air temperature T through a neural network A=NN(T, θ) and we optimize θ so that the generated solution matches some ice thickness reference. This reference is generated using the relation of the book from Cuffey and Paterson (2010).

Step 1: Parameter and glacier initialization

First we need to specify a list of RGI IDs of the glacier we want to work with. Specifying an RGI region is also possible. From these RGI IDs, we will look for the necessary files inside the workspace.

rgi_ids = ["RGI60-11.03638", "RGI60-11.01450", "RGI60-11.02346", "RGI60-07.00065"]
rgi_paths = get_rgi_paths()
Dict{String, String} with 56 entries:
  "RGI60-11.00897" => "per_glacier/RGI60-11/RGI60-11.00/RGI60-11.00897"
  "RGI60-08.00213" => "per_glacier/RGI60-08/RGI60-08.00/RGI60-08.00213"
  "RGI60-08.00147" => "per_glacier/RGI60-08/RGI60-08.00/RGI60-08.00147"
  "RGI60-11.01270" => "per_glacier/RGI60-11/RGI60-11.01/RGI60-11.01270"
  "RGI60-11.03646" => "per_glacier/RGI60-11/RGI60-11.03/RGI60-11.03646"
  "RGI60-11.03232" => "per_glacier/RGI60-11/RGI60-11.03/RGI60-11.03232"
  "RGI60-01.22174" => "per_glacier/RGI60-01/RGI60-01.22/RGI60-01.22174"
  "RGI60-07.00274" => "per_glacier/RGI60-07/RGI60-07.00/RGI60-07.00274"
  "RGI60-03.04207" => "per_glacier/RGI60-03/RGI60-03.04/RGI60-03.04207"
  "RGI60-04.04351" => "per_glacier/RGI60-04/RGI60-04.04/RGI60-04.04351"
  "RGI60-07.00065" => "per_glacier/RGI60-07/RGI60-07.00/RGI60-07.00065"
  "RGI60-11.02346" => "per_glacier/RGI60-11/RGI60-11.02/RGI60-11.02346"
  "RGI60-01.00570" => "per_glacier/RGI60-01/RGI60-01.00/RGI60-01.00570"
  "RGI60-11.01238" => "per_glacier/RGI60-11/RGI60-11.01/RGI60-11.01238"
  "RGI60-11.03005" => "per_glacier/RGI60-11/RGI60-11.03/RGI60-11.03005"
  "RGI60-08.00087" => "per_glacier/RGI60-08/RGI60-08.00/RGI60-08.00087"
  "RGI60-11.00787" => "per_glacier/RGI60-11/RGI60-11.00/RGI60-11.00787"
  "RGI60-08.00203" => "per_glacier/RGI60-08/RGI60-08.00/RGI60-08.00203"
  "RGI60-07.01193" => "per_glacier/RGI60-07/RGI60-07.01/RGI60-07.01193"
  ⋮                => ⋮

Define the time step for the simulation output and for the adjoint calculation. In this case, a month.

δt = 1/12
0.08333333333333333

Then we need to define the parameters of the simulation we want to perform. The arguments are very similar to the ones used in the forward simulation tutorial and for a complete explanation, the reader should refer to this tutorial. The main difference with the forward simulation tutorial here is that we need to specify the parameters for the functional inversion through the Hyperparameters and the UDEparameters. The Hyperparameters structure contains information about the optimization algorithm. The UDEparameters define how the Universal Differential Equation (UDE) is solved and how its gradient is computed.

params = Parameters(
    simulation = SimulationParameters(
        working_dir=working_dir,
        use_MB=false,
        use_velocities=true,
        tspan=(2010.0, 2015.0),
        step=δt,
        multiprocessing=true,
        workers=4,
        test_mode=false,
        rgi_paths=rgi_paths,
        gridScalingFactor=4), # Downscale the glacier grid to speed-up this example
    hyper = Hyperparameters(
        batch_size=length(rgi_ids), # We set batch size equals all datasize so we test gradient
        epochs=[15,10],
        optimizer=[ODINN.ADAM(0.01), ODINN.LBFGS(linesearch = ODINN.LineSearches.BackTracking(iterations = 5))]),
    physical = PhysicalParameters(
        minA = 8e-21,
        maxA = 8e-17),
    UDE = UDEparameters(
        optim_autoAD=ODINN.NoAD(),
        grad=ContinuousAdjoint(),
        optimization_method="AD+AD",
        target = :A),
    solver = Huginn.SolverParameters(
        step=δt,
        save_everystep=true,
        progress=true)
)
Sleipnir.Parameters{PhysicalParameters{Float64}, SimulationParameters{Int64, Float64, MeanDateVelocityMapping}, Hyperparameters{Float64, Int64}, SolverParameters{Float64, Int64}, UDEparameters{ContinuousAdjoint{Float64, Int64, DiscreteVJP}}, InversionParameters{Float64}}(PhysicalParameters{Float64}(900.0, 9.81, 1.0e-10, 1.0, 8.0e-17, 8.0e-21, 1.0, -25.0, 5.0e-18), SimulationParameters{Int64, Float64, MeanDateVelocityMapping}(false, true, true, true, false, false, (2010.0, 2015.0), 0.08333333333333333, true, 4, "/home/runner/work/ODINN.jl/ODINN.jl/demos", false, Dict("RGI60-11.00897" => "per_glacier/RGI60-11/RGI60-11.00/RGI60-11.00897", "RGI60-08.00213" => "per_glacier/RGI60-08/RGI60-08.00/RGI60-08.00213", "RGI60-08.00147" => "per_glacier/RGI60-08/RGI60-08.00/RGI60-08.00147", "RGI60-11.01270" => "per_glacier/RGI60-11/RGI60-11.01/RGI60-11.01270", "RGI60-11.03646" => "per_glacier/RGI60-11/RGI60-11.03/RGI60-11.03646", "RGI60-11.03232" => "per_glacier/RGI60-11/RGI60-11.03/RGI60-11.03232", "RGI60-01.22174" => "per_glacier/RGI60-01/RGI60-01.22/RGI60-01.22174", "RGI60-07.00274" => "per_glacier/RGI60-07/RGI60-07.00/RGI60-07.00274", "RGI60-03.04207" => "per_glacier/RGI60-03/RGI60-03.04/RGI60-03.04207", "RGI60-04.04351" => "per_glacier/RGI60-04/RGI60-04.04/RGI60-04.04351"…), "Farinotti19", MeanDateVelocityMapping(:nearest, 0.1643835616438356), 4), Hyperparameters{Float64, Int64}(1, 1, Float64[], Any[Optimisers.Adam(eta=0.01, beta=(0.9, 0.999), epsilon=1.0e-8), Optim.LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.BackTracking{Float64, Int64}, Returns{Nothing}}(10, LineSearches.InitialStatic{Float64}
  alpha: Float64 1.0
  scaled: Bool false
, LineSearches.BackTracking{Float64, Int64}
  c_1: Float64 0.0001
  ρ_hi: Float64 0.5
  ρ_lo: Float64 0.1
  iterations: Int64 5
  order: Int64 3
  maxstep: Float64 Inf
  cache: Nothing nothing
, nothing, Returns{Nothing}(nothing), Optim.Flat(), true)], 0.0, [15, 10], 4), SolverParameters{Float64, Int64}(OrdinaryDiffEqLowStorageRK.RDPK3Sp35{typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!), Static.False}(OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, static(false)), 1.0e-12, 0.08333333333333333, nothing, true, true, 10, 100000), UDEparameters{ContinuousAdjoint{Float64, Int64, DiscreteVJP}}(SciMLSensitivity.GaussAdjoint{0, true, Val{:central}, SciMLSensitivity.EnzymeVJP}(SciMLSensitivity.EnzymeVJP(0), false), SciMLBase.NoAD(), ContinuousAdjoint{Float64, Int64, DiscreteVJP}(DiscreteVJP(), OrdinaryDiffEqLowStorageRK.RDPK3Sp35{typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!), Static.False}(OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, static(false)), 1.0e-8, 1.0e-8, 0.08333333333333333, :Linear, 200), "AD+AD", LossH{L2Sum{Int64}}(L2Sum{Int64}(3)), :A), InversionParameters{Float64}([1.0], [0.0], [Inf], [1, 1], 0.001, 0.001, Optim.BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Optim.Flat}(LineSearches.InitialStatic{Float64}
  alpha: Float64 1.0
  scaled: Bool false
, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}
  delta: Float64 0.1
  sigma: Float64 0.9
  alphamax: Float64 Inf
  rho: Float64 5.0
  epsilon: Float64 1.0e-6
  gamma: Float64 0.66
  linesearchmax: Int64 50
  psi3: Float64 0.1
  display: Int64 0
  mayterminate: Base.RefValue{Bool}
  cache: Nothing nothing
  check_flatness: Bool false
, nothing, nothing, Optim.Flat())))

Then, we initialize those glaciers based on those RGI IDs and the parameters we previously specified.

glaciers = initialize_glaciers(rgi_ids, params)

Step 2: Defining a forward simulation as a synthetic ground truth

The next step is to generate a synthetic dataset using a forward simulation. This will generate a dataset with the ice thickness and surface velocities for each glacier at each time step. The dataset will be used to train the machine learning model. We define a synthetic law to generate the synthetic dataset. For this, we use some tabular data from Cuffey and Paterson (2010). The REPL shows that it maps the long term air temperature T to the creep coefficient A.

A_law = CuffeyPaterson()

The model is initialized using the Model constructor:

model = Huginn.Model(
    iceflow = SIA2Dmodel(params; A=A_law),
    mass_balance = TImodel1(params; DDF=6.0/1000.0, acc_factor=1.2/1000.0),
)

We define the time snapshots for transient inversion, i.e. the time steps at which we want to save the results, which will be used to compute the adjoint in reverse mode.

tstops = collect(2010:δt:2015)

prediction = Prediction(model, glaciers, params)
Prediction{Sleipnir.ModelCache{SIA2DCache{Float64, Int64, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}}, Nothing}}(Sleipnir.Model{SIA2Dmodel{Float64, Law{Array{Float64, 0}, Sleipnir.GenInputsAndApply{@NamedTuple{T::InpTemp}, Huginn.var"#34#36"{Polynomials.Polynomial{Float64, :x}}}, Huginn.var"#35#37", Nothing}, ConstantLaw{Array{Float64, 0}, Huginn.var"#9#10"}, ConstantLaw{Array{Float64, 0}, Huginn.var"#11#12"}, NullLaw, NullLaw}, TImodel1{Float64}, Nothing}(, , nothing), nothing, , Sleipnir.Parameters{PhysicalParameters{Float64}, SimulationParameters{Int64, Float64, MeanDateVelocityMapping}, Hyperparameters{Float64, Int64}, SolverParameters{Float64, Int64}, UDEparameters{ContinuousAdjoint{Float64, Int64, DiscreteVJP}}, InversionParameters{Float64}}(PhysicalParameters{Float64}(900.0, 9.81, 1.0e-10, 1.0, 8.0e-17, 8.0e-21, 1.0, -25.0, 5.0e-18), SimulationParameters{Int64, Float64, MeanDateVelocityMapping}(false, true, true, true, false, false, (2010.0, 2015.0), 0.08333333333333333, true, 4, "/home/runner/work/ODINN.jl/ODINN.jl/demos", false, Dict("RGI60-11.00897" => "per_glacier/RGI60-11/RGI60-11.00/RGI60-11.00897", "RGI60-08.00213" => "per_glacier/RGI60-08/RGI60-08.00/RGI60-08.00213", "RGI60-08.00147" => "per_glacier/RGI60-08/RGI60-08.00/RGI60-08.00147", "RGI60-11.01270" => "per_glacier/RGI60-11/RGI60-11.01/RGI60-11.01270", "RGI60-11.03646" => "per_glacier/RGI60-11/RGI60-11.03/RGI60-11.03646", "RGI60-11.03232" => "per_glacier/RGI60-11/RGI60-11.03/RGI60-11.03232", "RGI60-01.22174" => "per_glacier/RGI60-01/RGI60-01.22/RGI60-01.22174", "RGI60-07.00274" => "per_glacier/RGI60-07/RGI60-07.00/RGI60-07.00274", "RGI60-03.04207" => "per_glacier/RGI60-03/RGI60-03.04/RGI60-03.04207", "RGI60-04.04351" => "per_glacier/RGI60-04/RGI60-04.04/RGI60-04.04351"…), "Farinotti19", MeanDateVelocityMapping(:nearest, 0.1643835616438356), 4), Hyperparameters{Float64, Int64}(1, 1, Float64[], Any[Optimisers.Adam(eta=0.01, beta=(0.9, 0.999), epsilon=1.0e-8), Optim.LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.BackTracking{Float64, Int64}, Returns{Nothing}}(10, LineSearches.InitialStatic{Float64}
  alpha: Float64 1.0
  scaled: Bool false
, LineSearches.BackTracking{Float64, Int64}
  c_1: Float64 0.0001
  ρ_hi: Float64 0.5
  ρ_lo: Float64 0.1
  iterations: Int64 5
  order: Int64 3
  maxstep: Float64 Inf
  cache: Nothing nothing
, nothing, Returns{Nothing}(nothing), Optim.Flat(), true)], 0.0, [15, 10], 4), SolverParameters{Float64, Int64}(OrdinaryDiffEqLowStorageRK.RDPK3Sp35{typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!), Static.False}(OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, static(false)), 1.0e-12, 0.08333333333333333, nothing, true, true, 10, 100000), UDEparameters{ContinuousAdjoint{Float64, Int64, DiscreteVJP}}(SciMLSensitivity.GaussAdjoint{0, true, Val{:central}, SciMLSensitivity.EnzymeVJP}(SciMLSensitivity.EnzymeVJP(0), false), SciMLBase.NoAD(), ContinuousAdjoint{Float64, Int64, DiscreteVJP}(DiscreteVJP(), OrdinaryDiffEqLowStorageRK.RDPK3Sp35{typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!), Static.False}(OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, static(false)), 1.0e-8, 1.0e-8, 0.08333333333333333, :Linear, 200), "AD+AD", LossH{L2Sum{Int64}}(L2Sum{Int64}(3)), :A), InversionParameters{Float64}([1.0], [0.0], [Inf], [1, 1], 0.001, 0.001, Optim.BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Optim.Flat}(LineSearches.InitialStatic{Float64}
  alpha: Float64 1.0
  scaled: Bool false
, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}
  delta: Float64 0.1
  sigma: Float64 0.9
  alphamax: Float64 Inf
  rho: Float64 5.0
  epsilon: Float64 1.0e-6
  gamma: Float64 0.66
  linesearchmax: Int64 50
  psi3: Float64 0.1
  display: Int64 0
  mayterminate: Base.RefValue{Bool}
  cache: Nothing nothing
  check_flatness: Bool false
, nothing, nothing, Optim.Flat()))), Sleipnir.Results[])

We generate the synthetic dataset using the forward simulation. This will generate a dataset with the ice thickness and surface velocities for each glacier at each time step. The dataset will be used to train the machine learning model. This will run under the hood a Prediction using Huginn.jl.

glaciers = generate_ground_truth(glaciers, params, model, tstops)

The results of this simulation are stored in the thicknessData field of each glacier.

Step 3: Model specification to perform a functional inversion

After this forward simulation, we define a new iceflow model to be ready for the inversions. The first step is to define a simple neural network that takes as input a scalar and returns a scalar.

nn_model = NeuralNetwork(params)

Then we define a law that uses this neural network to map the long term air temperature T to the creep coefficient A. ODINN comes with a set of already defined laws. Only a few of them support functional inversion as the computation of the gradient needs to be carefully handled. More information about these laws can be found in the laws tutorial.

A_law = LawA(nn_model, params)

Then we define an iceflow and ODINN tells us how the law is used in the iceflow equation.

iceflow = SIA2Dmodel(params; A=A_law)

Finally we define the model which needs to know the iceflow and mass balance models, and in comparison to Huginn, there is a third argument regressors. This regressors argument tells how each regressor relates into the SIA. Although we already defined this in the iceflow model, this definition is mandatory for technical reasons. This argument will probably disappear in the future once the code becomes more mature. It must match how the laws are defined in the iceflow model.

model = Model(
    iceflow = iceflow,
    mass_balance = TImodel1(params; DDF=6.0/1000.0, acc_factor=1.2/1000.0),
    regressors = (; A=nn_model)
)

Step 4: Train a Universal Differential Equation via a functional inversion

The next step is to specify the type of simulation we want to perform. In this case, we will use a FunctionalInversion simulation, which will use the synthetic dataset generated in the previous step to train a Universal Differential Equation (UDE) model.

functional_inversion = FunctionalInversion(model, glaciers, params)
FunctionalInversion{Sleipnir.Model{SIA2Dmodel{Float64, Law{Array{Float64, 0}, Sleipnir.GenInputsAndApply{@NamedTuple{T::InpTemp}, ODINN.var"#220#223"{Lux.StatefulLuxLayer{Static.True, Lux.Chain{@NamedTuple{layer_1::Lux.Dense{ODINN.var"#107#111", Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{ODINN.var"#108#112", Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Lux.Dense{ODINN.var"#109#113", Int64, Int64, Nothing, Nothing, Static.True}, layer_4::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}, layer_4::@NamedTuple{}}}, Sleipnir.Parameters{PhysicalParameters{Float64}, SimulationParameters{Int64, Float64, MeanDateVelocityMapping}, Hyperparameters{Float64, Int64}, SolverParameters{Float64, Int64}, UDEparameters{ContinuousAdjoint{Float64, Int64, DiscreteVJP}}, InversionParameters{Float64}}}}, ODINN.var"#222#225", Nothing}, ConstantLaw{Array{Float64, 0}, Huginn.var"#9#10"}, ConstantLaw{Array{Float64, 0}, Huginn.var"#11#12"}, NullLaw, NullLaw}, TImodel1{Float64}, ODINN.MachineLearning{NeuralNetwork{Lux.Chain{@NamedTuple{layer_1::Lux.Dense{ODINN.var"#107#111", Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{ODINN.var"#108#112", Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Lux.Dense{ODINN.var"#109#113", Int64, Int64, Nothing, Nothing, Static.True}, layer_4::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, ComponentArrays.ComponentVector{Float64, Vector{Float64}, Tuple{ComponentArrays.Axis{(θ = ViewAxis(1:83, Axis(layer_1 = ViewAxis(1:6, Axis(weight = ViewAxis(1:3, ShapedAxis((3, 1))), bias = ViewAxis(4:6, Shaped1DAxis((3,))))), layer_2 = ViewAxis(7:46, Axis(weight = ViewAxis(1:30, ShapedAxis((10, 3))), bias = ViewAxis(31:40, Shaped1DAxis((10,))))), layer_3 = ViewAxis(47:79, Axis(weight = ViewAxis(1:30, ShapedAxis((3, 10))), bias = ViewAxis(31:33, Shaped1DAxis((3,))))), layer_4 = ViewAxis(80:83, Axis(weight = ViewAxis(1:3, ShapedAxis((1, 3))), bias = ViewAxis(4:4, Shaped1DAxis((1,))))))),)}}}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}, layer_4::@NamedTuple{}}}, ODINN.emptyMLmodel, ODINN.emptyMLmodel, ODINN.emptyMLmodel, ODINN.emptyMLmodel, SIA2D_A_target, ComponentArrays.ComponentVector{Float64, Vector{Float64}, Tuple{ComponentArrays.Axis{(A = ViewAxis(1:83, Axis(layer_1 = ViewAxis(1:6, Axis(weight = ViewAxis(1:3, ShapedAxis((3, 1))), bias = ViewAxis(4:6, Shaped1DAxis((3,))))), layer_2 = ViewAxis(7:46, Axis(weight = ViewAxis(1:30, ShapedAxis((10, 3))), bias = ViewAxis(31:40, Shaped1DAxis((10,))))), layer_3 = ViewAxis(47:79, Axis(weight = ViewAxis(1:30, ShapedAxis((3, 10))), bias = ViewAxis(31:33, Shaped1DAxis((3,))))), layer_4 = ViewAxis(80:83, Axis(weight = ViewAxis(1:3, ShapedAxis((1, 3))), bias = ViewAxis(4:4, Shaped1DAxis((1,))))))),)}}}}}, Sleipnir.ModelCache{SIA2DCache{Float64, Int64, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}, Array{Float64, 0}}, Nothing}, Glacier2D{Float64, Int64, Climate2D{Rasters.RasterStack{(:prcp, :temp, :gradient), @NamedTuple{prcp::Float64, temp::Float64, gradient::Float64}, 1, @NamedTuple{prcp::Vector{Float64}, temp::Vector{Float64}, gradient::Vector{Float64}}, Tuple{DimensionalData.Dimensions.Ti{DimensionalData.Dimensions.Lookups.Sampled{Dates.DateTime, Vector{Dates.DateTime}, DimensionalData.Dimensions.Lookups.ForwardOrdered, DimensionalData.Dimensions.Lookups.Irregular{Tuple{Nothing, Nothing}}, DimensionalData.Dimensions.Lookups.Points, DimensionalData.Dimensions.Lookups.Metadata{Rasters.NCDsource, Dict{String, Any}}}}}, Tuple{}, @NamedTuple{prcp::Tuple{DimensionalData.Dimensions.Ti{Colon}}, temp::Tuple{DimensionalData.Dimensions.Ti{Colon}}, gradient::Tuple{DimensionalData.Dimensions.Ti{Colon}}}, DimensionalData.Dimensions.Lookups.Metadata{Rasters.NCDsource, Dict{String, Any}}, @NamedTuple{prcp::DimensionalData.Dimensions.Lookups.Metadata{Rasters.NCDsource, Dict{String, Any}}, temp::DimensionalData.Dimensions.Lookups.Metadata{Rasters.NCDsource, Dict{String, Any}}, gradient::DimensionalData.Dimensions.Lookups.Metadata{Rasters.NCDsource, Dict{String, Any}}}, Nothing}, Rasters.RasterStack{(:prcp, :temp, :gradient), @NamedTuple{prcp::Float64, temp::Float64, gradient::Float64}, 1, @NamedTuple{prcp::Vector{Float64}, temp::Vector{Float64}, gradient::Vector{Float64}}, Tuple{DimensionalData.Dimensions.Ti{DimensionalData.Dimensions.Lookups.Sampled{Dates.DateTime, Vector{Dates.DateTime}, DimensionalData.Dimensions.Lookups.ForwardOrdered, DimensionalData.Dimensions.Lookups.Irregular{Tuple{Nothing, Nothing}}, DimensionalData.Dimensions.Lookups.Points, DimensionalData.Dimensions.Lookups.Metadata{Rasters.NCDsource, Dict{String, Any}}}}}, Tuple{}, @NamedTuple{prcp::Tuple{DimensionalData.Dimensions.Ti{Colon}}, temp::Tuple{DimensionalData.Dimensions.Ti{Colon}}, gradient::Tuple{DimensionalData.Dimensions.Ti{Colon}}}, DimensionalData.Dimensions.Lookups.Metadata{Rasters.NCDsource, Dict{String, Any}}, @NamedTuple{prcp::DimensionalData.Dimensions.Lookups.Metadata{Rasters.NCDsource, Dict{String, Any}}, temp::DimensionalData.Dimensions.Lookups.Metadata{Rasters.NCDsource, Dict{String, Any}}, gradient::DimensionalData.Dimensions.Lookups.Metadata{Rasters.NCDsource, Dict{String, Any}}}, Nothing}, Sleipnir.ClimateStep{Float64}, Climate2Dstep{Float64}, Float64}, ThicknessData{Float64}, SurfaceVelocityData{Float64}}, Results{Sleipnir.Results{Float64, Int64}, TrainingStats{Float64, Int64}}}(Sleipnir.Model{SIA2Dmodel{Float64, Law{Array{Float64, 0}, Sleipnir.GenInputsAndApply{@NamedTuple{T::InpTemp}, ODINN.var"#220#223"{Lux.StatefulLuxLayer{Static.True, Lux.Chain{@NamedTuple{layer_1::Lux.Dense{ODINN.var"#107#111", Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{ODINN.var"#108#112", Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Lux.Dense{ODINN.var"#109#113", Int64, Int64, Nothing, Nothing, Static.True}, layer_4::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}, layer_4::@NamedTuple{}}}, Sleipnir.Parameters{PhysicalParameters{Float64}, SimulationParameters{Int64, Float64, MeanDateVelocityMapping}, Hyperparameters{Float64, Int64}, SolverParameters{Float64, Int64}, UDEparameters{ContinuousAdjoint{Float64, Int64, DiscreteVJP}}, InversionParameters{Float64}}}}, ODINN.var"#222#225", Nothing}, ConstantLaw{Array{Float64, 0}, Huginn.var"#9#10"}, ConstantLaw{Array{Float64, 0}, Huginn.var"#11#12"}, NullLaw, NullLaw}, TImodel1{Float64}, ODINN.MachineLearning{NeuralNetwork{Lux.Chain{@NamedTuple{layer_1::Lux.Dense{ODINN.var"#107#111", Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{ODINN.var"#108#112", Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Lux.Dense{ODINN.var"#109#113", Int64, Int64, Nothing, Nothing, Static.True}, layer_4::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, ComponentArrays.ComponentVector{Float64, Vector{Float64}, Tuple{ComponentArrays.Axis{(θ = ViewAxis(1:83, Axis(layer_1 = ViewAxis(1:6, Axis(weight = ViewAxis(1:3, ShapedAxis((3, 1))), bias = ViewAxis(4:6, Shaped1DAxis((3,))))), layer_2 = ViewAxis(7:46, Axis(weight = ViewAxis(1:30, ShapedAxis((10, 3))), bias = ViewAxis(31:40, Shaped1DAxis((10,))))), layer_3 = ViewAxis(47:79, Axis(weight = ViewAxis(1:30, ShapedAxis((3, 10))), bias = ViewAxis(31:33, Shaped1DAxis((3,))))), layer_4 = ViewAxis(80:83, Axis(weight = ViewAxis(1:3, ShapedAxis((1, 3))), bias = ViewAxis(4:4, Shaped1DAxis((1,))))))),)}}}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}, layer_4::@NamedTuple{}}}, ODINN.emptyMLmodel, ODINN.emptyMLmodel, ODINN.emptyMLmodel, ODINN.emptyMLmodel, SIA2D_A_target, ComponentArrays.ComponentVector{Float64, Vector{Float64}, Tuple{ComponentArrays.Axis{(A = ViewAxis(1:83, Axis(layer_1 = ViewAxis(1:6, Axis(weight = ViewAxis(1:3, ShapedAxis((3, 1))), bias = ViewAxis(4:6, Shaped1DAxis((3,))))), layer_2 = ViewAxis(7:46, Axis(weight = ViewAxis(1:30, ShapedAxis((10, 3))), bias = ViewAxis(31:40, Shaped1DAxis((10,))))), layer_3 = ViewAxis(47:79, Axis(weight = ViewAxis(1:30, ShapedAxis((3, 10))), bias = ViewAxis(31:33, Shaped1DAxis((3,))))), layer_4 = ViewAxis(80:83, Axis(weight = ViewAxis(1:3, ShapedAxis((1, 3))), bias = ViewAxis(4:4, Shaped1DAxis((1,))))))),)}}}}}(, , ), nothing, , Sleipnir.Parameters{PhysicalParameters{Float64}, SimulationParameters{Int64, Float64, MeanDateVelocityMapping}, Hyperparameters{Float64, Int64}, SolverParameters{Float64, Int64}, UDEparameters{ContinuousAdjoint{Float64, Int64, DiscreteVJP}}, InversionParameters{Float64}}(PhysicalParameters{Float64}(900.0, 9.81, 1.0e-10, 1.0, 8.0e-17, 8.0e-21, 1.0, -25.0, 5.0e-18), SimulationParameters{Int64, Float64, MeanDateVelocityMapping}(false, true, true, true, false, false, (2010.0, 2015.0), 0.08333333333333333, true, 4, "/home/runner/work/ODINN.jl/ODINN.jl/demos", false, Dict("RGI60-11.00897" => "per_glacier/RGI60-11/RGI60-11.00/RGI60-11.00897", "RGI60-08.00213" => "per_glacier/RGI60-08/RGI60-08.00/RGI60-08.00213", "RGI60-08.00147" => "per_glacier/RGI60-08/RGI60-08.00/RGI60-08.00147", "RGI60-11.01270" => "per_glacier/RGI60-11/RGI60-11.01/RGI60-11.01270", "RGI60-11.03646" => "per_glacier/RGI60-11/RGI60-11.03/RGI60-11.03646", "RGI60-11.03232" => "per_glacier/RGI60-11/RGI60-11.03/RGI60-11.03232", "RGI60-01.22174" => "per_glacier/RGI60-01/RGI60-01.22/RGI60-01.22174", "RGI60-07.00274" => "per_glacier/RGI60-07/RGI60-07.00/RGI60-07.00274", "RGI60-03.04207" => "per_glacier/RGI60-03/RGI60-03.04/RGI60-03.04207", "RGI60-04.04351" => "per_glacier/RGI60-04/RGI60-04.04/RGI60-04.04351"…), "Farinotti19", MeanDateVelocityMapping(:nearest, 0.1643835616438356), 4), Hyperparameters{Float64, Int64}(1, 1, Float64[], Any[Optimisers.Adam(eta=0.01, beta=(0.9, 0.999), epsilon=1.0e-8), Optim.LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.BackTracking{Float64, Int64}, Returns{Nothing}}(10, LineSearches.InitialStatic{Float64}
  alpha: Float64 1.0
  scaled: Bool false
, LineSearches.BackTracking{Float64, Int64}
  c_1: Float64 0.0001
  ρ_hi: Float64 0.5
  ρ_lo: Float64 0.1
  iterations: Int64 5
  order: Int64 3
  maxstep: Float64 Inf
  cache: Nothing nothing
, nothing, Returns{Nothing}(nothing), Optim.Flat(), true)], 0.0, [15, 10], 4), SolverParameters{Float64, Int64}(OrdinaryDiffEqLowStorageRK.RDPK3Sp35{typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!), Static.False}(OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, static(false)), 1.0e-12, 0.08333333333333333, nothing, true, true, 10, 100000), UDEparameters{ContinuousAdjoint{Float64, Int64, DiscreteVJP}}(SciMLSensitivity.GaussAdjoint{0, true, Val{:central}, SciMLSensitivity.EnzymeVJP}(SciMLSensitivity.EnzymeVJP(0), false), SciMLBase.NoAD(), ContinuousAdjoint{Float64, Int64, DiscreteVJP}(DiscreteVJP(), OrdinaryDiffEqLowStorageRK.RDPK3Sp35{typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!), Static.False}(OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, static(false)), 1.0e-8, 1.0e-8, 0.08333333333333333, :Linear, 200), "AD+AD", LossH{L2Sum{Int64}}(L2Sum{Int64}(3)), :A), InversionParameters{Float64}([1.0], [0.0], [Inf], [1, 1], 0.001, 0.001, Optim.BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Nothing, Optim.Flat}(LineSearches.InitialStatic{Float64}
  alpha: Float64 1.0
  scaled: Bool false
, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}
  delta: Float64 0.1
  sigma: Float64 0.9
  alphamax: Float64 Inf
  rho: Float64 5.0
  epsilon: Float64 1.0e-6
  gamma: Float64 0.66
  linesearchmax: Int64 50
  psi3: Float64 0.1
  display: Int64 0
  mayterminate: Base.RefValue{Bool}
  cache: Nothing nothing
  check_flatness: Bool false
, nothing, nothing, Optim.Flat()))), Results{Sleipnir.Results{Float64, Int64}, TrainingStats{Float64, Int64}}(Sleipnir.Results{Float64, Int64}[], TrainingStats{Float64, Int64}(nothing, Float64[], 0, nothing, ComponentArrays.ComponentVector[], ComponentArrays.ComponentVector[], Dates.DateTime("0000-01-01T00:00:00"))))

And finally, we just run the simulation. This will run the adjoint method to compute the gradients and then use the ADAM optimizer to train the UDE model.

run!(functional_inversion)
Running training of UDE...

[ Info: Training with ADAM optimizer
[ Info: Training with custom ContinuousAdjoint{Float64, Int64, DiscreteVJP} method
Iteration: [    1 /    25]     Loss:7.94252e+01
Iteration: [    2 /    25]     Loss:6.57928e+01     Improvement: -17.16 %
Iteration: [    3 /    25]     Loss:5.49244e+01     Improvement: -16.52 %
Iteration: [    4 /    25]     Loss:4.63710e+01     Improvement: -15.57 %
Iteration: [    5 /    25]     Loss:3.96132e+01     Improvement: -14.57 %
Iteration: [    6 /    25]     Loss:3.41950e+01     Improvement: -13.68 %
Iteration: [    7 /    25]     Loss:2.97769e+01     Improvement: -12.92 %
Iteration: [    8 /    25]     Loss:2.61246e+01     Improvement: -12.27 %
Iteration: [    9 /    25]     Loss:2.30760e+01     Improvement: -11.67 %
Iteration: [   10 /    25]     Loss:2.05147e+01     Improvement: -11.10 %
Iteration: [   11 /    25]     Loss:1.83530e+01     Improvement: -10.54 %
Iteration: [   12 /    25]     Loss:1.65219e+01     Improvement: -9.98 %
Iteration: [   13 /    25]     Loss:1.49655e+01     Improvement: -9.42 %
Iteration: [   14 /    25]     Loss:1.36379e+01     Improvement: -8.87 %
Iteration: [   15 /    25]     Loss:1.25008e+01     Improvement: -8.34 %
Iteration: [   16 /    25]     Loss:1.25008e+01     Improvement: 0.00 %
[ Info: Training with BFGS optimizer
[ Info: Training with custom ContinuousAdjoint{Float64, Int64, DiscreteVJP} method
Iteration: [   17 /    25]     Loss:1.25008e+01     Improvement: 0.00 %
Iteration: [   18 /    25]     Loss:1.04657e+01     Improvement: -16.28 %

This page was generated using Literate.jl.