Extending ODINN

This page is the consolidated guide for developers who want to add new physics, models, or algorithms to the ODINN ecosystem. Each section corresponds to one extension point and describes the minimum interface to implement, plus pointers to deeper material.

Add a new iceflow model

New iceflow models subtype IceflowModel (defined in Sleipnir). SIAmodel is an intermediate abstract type for Shallow Ice Approximation variants; a model based on a different physical approximation — Shallow Shelf Approximation, DIVA, etc. — sits directly under IceflowModel.

Before reading this section, it helps to understand how the existing SIA2D model is wired into OrdinaryDiffEq.jl — see the execution chain diagram in the Huginn package page.

The interface builds up in three layers depending on what you need:

Layer 1 — Forward simulation (run ice thickness evolution, no gradients):

Layer 2 — Inversion and adjoint differentiation (needed if you want to use this model with Inversion or UDE training):

Layer 3 — Surface velocity diagnostics (optional):

Example skeleton — SSA2Dmodel

Five pieces are needed: the model struct, the cache struct, cache initialization, the PDE kernel, and the batch_iceflow_PDE! override.

using Huginn, Sleipnir

# Type hierarchy:
#   IceflowModel (abstract, Sleipnir)
#     ├── SIAmodel (abstract) → SIA2Dmodel
#     └── SSA2Dmodel  ← new model directly under IceflowModel

# ── Model type: holds laws and configuration ──────────────────────────────
struct SSA2Dmodel <: IceflowModel
    viscosity_law::Any   # law for effective viscosity η
    friction_law::Any    # law for basal friction coefficient
    # ... (add all law fields required by the SSA kernel)
end

# ── Cache type: pre-allocated arrays, reused at every ODE step ────────────
# For a complete reference on what to pre-allocate, see SIA2DCache in SIA2D.jl
mutable struct SSA2DCache
    Ux::Matrix{Float64}   # x-velocity
    Uy::Matrix{Float64}   # y-velocity
    # ... (staggered-grid fields, stress tensors, law caches, MB fields, etc.)
    glacier_idx::Int
end

# ── Cache initializer: called once per glacier before the ODE solve ───────
function Sleipnir.init_cache(model::SSA2Dmodel, simulation, glacier_idx::Int, θ)
    g = simulation.glaciers[glacier_idx]
    nx, ny = g.nx, g.ny
    return SSA2DCache(zeros(nx, ny), zeros(nx, ny), glacier_idx)
end

# ── Required: law application (non-callback laws applied inside ODE step) ─
function Huginn.apply_all_non_callback_laws!(model::SSA2Dmodel, cache::SSA2DCache,
        simulation, glacier_idx, t, θ)
    apply_law!(model.viscosity_law, cache.η, simulation, glacier_idx, t, θ)
    # ... apply all non-callback laws for this model
end

# ── Required: law application (callback laws applied at discrete steps) ───
function Huginn.apply_all_callback_laws!(model::SSA2Dmodel, cache::SSA2DCache,
        simulation, glacier_idx, t, θ)
    # apply laws that fire at callback frequency (e.g. sliding)
end

# ── Required for forward stub (inversion override lives in ODINN.jl) ──────
function Huginn.precompute_all_VJPs_laws!(model::SSA2Dmodel, cache::SSA2DCache,
        simulation::Prediction, glacier_idx, t, θ)
    nothing
end

# ── Callbacks: periodic law updates ───────────────────────────────────────
function Huginn.build_callback(model::SSA2Dmodel, cache::SSA2DCache,
        glacier_idx, tspan)
    return CallbackSet()
end

# ── PDE kernel: analogous to SIA2D! — keep θ for AD compatibility ─────────
function SSA2D!(dU, U, simulation, t, θ)
    # write ∂U/∂t into dU using the SSA stress balance and mass continuity ...
end

# ── ODE adapter: drops θ to match ODEProblem's f(du,u,p,t) interface ─────
function SSA2D_PDE!(dU, U, simulation, t)
    SSA2D!(dU, U, simulation, t, nothing)
end

# ── Entry point: override batch_iceflow_PDE! to wire in SSA2D_PDE! ────────
function Huginn.batch_iceflow_PDE!(glacier_idx::Int,
        simulation::Prediction{<:Sleipnir.Model{SSA2Dmodel}})
    params = simulation.parameters
    simulation.cache = Sleipnir.init_cache(simulation.model, simulation, glacier_idx, nothing)
    tstops = Huginn.define_callback_steps(params.simulation.tspan, params.solver.step)
    cb = build_callback(simulation.model.iceflow, simulation.cache.iceflow,
        glacier_idx, params.simulation.tspan)
    return Huginn.simulate_iceflow_PDE!(simulation, cb, SSA2D_PDE!, tstops)
end

Add a new mass balance model

New mass balance models subtype MBmodel (defined in Muninn). The MB callback fires every step_MB and runs three steps in sequence — only compute_MB is model-specific:

MB callback (every step_MB):
  ├── MB_timestep!(cache, model, glacier, step, t, glacier_idx)  # writes the MB into cache.iceflow.MB
  │     └── compute_MB(mb_model, climate_2D_step, step)     ← implement this for your model
  ├── apply_MB_mask!(H, cache.iceflow)                      # applies the MB to the ice thickness H, clipping to avoid negative thickness
  └── push!(cache.iceflow.MB_history, copy(cache.iceflow.MB))   # records the MB snapshot

Minimum to implement:

using Muninn

struct MyMBmodel <: MBmodel
    # your fields
end

# Required: compute the distributed MB for one time step
function Muninn.compute_MB(model::MyMBmodel, climate_step::Climate2Dstep,
        step::AbstractFloat)
    # climate_step — gridded climate fields (temp, prcp, PDD, etc.)
    # step — fractional year length of this timestep
    # return a (nx, ny) matrix in m w.e.
end

Optional dispatch hooks (all have sensible defaults in Muninn — override only what differs):

Muninn.requires_dynamic_topography(::MyMBmodel) = false   # true if model uses slope/aspect
Muninn.topography_window_m(::MyMBmodel) = 200.0   # DEM smoothing radius (m)
Muninn.mb_inputs(::MyMBmodel) = (;)     # extra NamedTuple inputs
Muninn.required_climate_data_source(::MyMBmodel) = nothing # :ERA5 or :W5E5
Muninn.get_temp_bias(::MyMBmodel) = 0.0     # temperature offset (°C)

Pass your model to Model(; iceflow = iceflow_model, mass_balance = MyMBmodel(...), regressors = ...) as usual. Note that Model itself is defined in Sleipnir; ODINN only extends it (through _construct_Model) to build the TrainableComponents when regressors are provided.

Add a new iceflow law

Laws are the primary mechanism for injecting custom or learnable physics into the iceflow solver. A Law wraps a computation — pure physics or a neural network — and is called at each ODE step (or at a fixed callback frequency).

Where to add new law types:

• Learnable law (wraps a regressor, used in UDE training): add to ODINN.jl/src/laws/Laws.jl
• Non-learnable law (pure physics, no neural network): add to Huginn.jl/src/laws/Laws.jl

Minimal example — a non-learnable diffusivity law:

using Sleipnir

struct MyDiffusivityLaw <: AbstractLaw{Matrix{Float64}}
    name::Symbol
    inputs::NamedTuple
    f!::Function
    init_cache::Function
    callback_freq::Union{Nothing, Real}
end

function MyDiffusivityLaw(; inputs = (;))
    MyDiffusivityLaw(
        :MyD,
        inputs,
        # f! receives (cache, inp, θ): inp is the NamedTuple of resolved inputs, θ holds NN params
        (cache,
            inp,
            θ) -> @. cache.output = inp.H ^ 3,
        (model, glacier) -> MatrixCache(glacier.nx, glacier.ny),
        nothing  # no callback; apply at every ODE step
    )
end

See the Laws tutorial for complete worked examples (learnable and non-learnable), the Laws inputs tutorial for implementing custom AbstractInput types, and the Laws VJP tutorial for customizing adjoints for performance-sensitive laws.

For the conceptual overview of how Law binds inputs and a regressor to a target component, see the Inversions page.

Add a new loss function

A loss function measures the mismatch between the model's predicted state (ice thickness, surface velocity, etc.) and observations. For most of the losses, the metric itself is a simple loss (AbstractSimpleLoss, like the built-in L2Sum and LogSum); the composites LossH, LossV, LossHV (subtypes of AbstractLoss) then apply that metric to ice thickness and/or velocity. To add a new metric (e.g. a mean absolute error), subtype AbstractSimpleLoss. It needs a distance field — the composite uses it to build the in-glacier mask — and a loss method returning a scalar:

Beyond the per-timestep composites above, ODINN also provides time-aggregated losses (LossDhdt, LossAvgV, subtypes of TimeAggregatedLoss) which compare quantities integrated over the simulation window rather than pointwise in time, MultiLoss to combine several losses with weights, and a family of regularization terms (TikhonovRegularization, InitialThicknessRegularization, VelocityRegularization, RheologyRegularization, DiffusivityRegularization). Not all of these decompose into a simple-loss metric — LossDhdt, for instance, defines its own aggregation — so use them as templates when your new loss does not fit the AbstractSimpleLoss shape.

using ODINN

struct MyLoss <: AbstractSimpleLoss
    distance::Int
end
MyLoss(; distance = 3) = MyLoss(distance)

# a = prediction, b = reference (both (nx, ny)); mask is TRUE for valid in-glacier
# pixels; normalization is a scalar divisor. Must return a scalar.
function ODINN.loss(::MyLoss, a::Matrix, b::Matrix, mask::BitMatrix, normalization)
    return sum(abs.(a[mask] .- b[mask])) / normalization
end

Select it by wrapping it in a thickness/velocity loss and passing it through UDEparameters (there is no loss keyword on Inversion):

params = Parameters(
# …,
    UDE = UDEparameters(empirical_loss_function = LossH(loss = MyLoss()))   # LossV / LossHV for velocity
)

Do you also need backward_loss?

backward_loss returns ∂L/∂a (same shape as a), zero outside the mask. It is only called by ODINN's manual adjoint methods (DiscreteAdjoint and ContinuousAdjoint). With SciMLSensitivityAdjoint (configured via UDEparameters(grad = SciMLSensitivityAdjoint(), optim_autoAD = Optimization.AutoZygote())), Zygote differentiates through the loss automatically and backward_loss is never called. You only need it for the manual adjoints:

function ODINN.backward_loss(::MyLoss, a::Matrix, b::Matrix, mask::BitMatrix, normalization)
    # ∂/∂a of sum|a − b|, restricted to valid pixels
    d = zero(a)
    d[mask] = sign.(a[mask] .- b[mask])
    return d ./ normalization
end

See Sensitivity analysis for a guide to choosing between adjoint methods.

Add a new inversion target

When do you need a custom target? Only if you use ODINN's manual adjoint methods (ContinuousAdjoint or DiscreteAdjoint). Those methods require an explicit AbstractSIA2DTarget that hand-codes how your quantity enters the SIA2D diffusivity Jacobians. Currently implemented targets cover A (Glen flow rate factor) and D (diffusivity): SIA2D_A_target, SIA2D_D_target, SIA2D_D_hybrid_target.

If you use SciMLSensitivityAdjoint instead, no custom target is needed — Zygote + SciMLSensitivity differentiate through the full ODE automatically. This means parameters like the basal sliding coefficient C can already be inverted today via SciMLSensitivityAdjoint, simply by adding a C law to SIA2Dmodel and registering it in TrainableComponents — no new target code required.

To add a manual adjoint target for a new quantity, subtype AbstractSIA2DTarget and implement:

• Diffusivity(target; H̄, ∇S, θ, ...) — the full diffusivity expression, including the contribution of the target quantity
• ∂Diffusivity∂H, ∂Diffusivity∂∇H, ∂Diffusivity∂θ — staggered-grid derivatives used by the adjoint
• Optionally Velocityꜛ and its derivatives if you fit to surface velocity observations

The existing targets in src/models/target/ are the reference: SIA2D_A_target is the simplest (one scalar field), SIA2D_D_hybrid_target the most complex (combines A and D). Copy the nearest analogue and adapt the PDE terms.