Numerical implementation

This page describes how Ribasim’s Julia core implements the numerical methods discussed in numerical considerations and solves the equations defined in equations. It focuses on the SciML ecosystem integration, the automatic differentiation pipeline, and the key implementation patterns in core/src/.

1 SciML ecosystem overview

Ribasim builds on several packages from the SciML ecosystem:

Package	Role
OrdinaryDiffEq.jl subpackages	Time integration (default solver: NordsieckBDF)
DifferentiationInterface.jl	Unified API for Jacobian computation
ForwardDiff.jl	Forward-mode automatic differentiation via dual numbers
SparseConnectivityTracer.jl	Sparsity pattern detection
SparseMatrixColorings.jl	Graph coloring for efficient sparse Jacobian evaluation
LinearSolve.jl	Linear solver interface (KLU factorization)
DataInterpolations.jl	Interpolation of forcing data and other tabular inputs (e.g. Basin profiles, rating curves, control look-up functions)

2 ODE problem setup

The ODE problem is assembled in core/src/model.jl. The key steps are:

Build the Parameters struct containing all model data
Construct an in-place ODEFunction wrapping water_balance!, with the Jacobian (prototype) and time derivative set up by get_diff_eval (core/src/solve.jl); by default the Jacobian prototype is the lazy RibasimJacobian operator
Wrap it in an ODEProblem and create the integrator with init, registering all callbacks

OrdinaryDiffEq’s specialization level is also chosen here: FullSpecialize produces the fastest code but compiles on every type change, while NoSpecialize reduces compilation time at the cost of some runtime performance.

The simulation is then advanced with step!(integrator) in a loop until t_end.

3 The right-hand side: `water_balance!`

The RHS function water_balance! in core/src/solve.jl has the in-place signature water_balance!(du, u, p, t) expected by OrdinaryDiffEq. It writes the time derivatives of all cumulative flow states into the pre-allocated output vector du (the same length as the state vector u); no allocations or return value are involved. Its execution follows a fixed order:

Cache check — check_new_input! determines whether time-dependent or state-dependent cached quantities need recomputation
Zero du — every entry is reset before any flow contributions are accumulated into it
Basin properties — compute current storage, level, area, and the low storage factor from the cumulative flows in u
Vertical fluxes — precipitation and drainage are state-independent and added directly, while evaporation and infiltration are state dependent; this asymmetry is what makes the outgoing fluxes contribute to the Jacobian while the incoming ones do not
Horizontal flows — formulate_flows! for each node type (Pump, Outlet, LinearResistance, ManningResistance, TabulatedRatingCurve, UserDemand, …)
Continuous control — evaluate continuous control rules that depend on current flows and levels
Controlled flows — flows whose parameters are set by continuous control
PID control — evaluate PID control (integral/proportional/derivative)
PID-controlled flows — flows whose parameters are set by PID control

Note

The state vector \(\mathbf{u}\) contains cumulative flows, not instantaneous flows or storages. See why this formulation for the rationale.

3.1 Parameter structure

The Parameters struct (defined in core/src/parameter.jl) is split into four components that play distinct roles during Jacobian computation:

Component	Contents
`p_independent`	All node data, graph, interpolations, pre-allocated buffers
`current_basin_properties`	Cached values that depend on both state and time (current storage, level, area, low storage factor)
`time_dependent_cache`	Cached values that depend only on time (forcing interpolations)
`p_mutable`	Flags controlling cache validity (`new_time_dependent_cache`), Jacobian refreshing (`refresh_jac`) and AD (`ad_active`)

This split is driven by AD requirements — see Jacobian computation below.

4 Implicit time integration

Water systems can be stiff — some processes change on timescales of seconds (a pump switching on), while others evolve over days (a reservoir slowly filling). Explicit methods like forward Euler would need extremely small timesteps to remain stable, making them impractical. Implicit methods avoid this by solving an equation at each step that accounts for the system’s coupled dynamics, allowing much larger steps.

The default solver is NordsieckBDF (backward differentiation formula), an implicit multi-step method from OrdinaryDiffEq.jl well-suited for stiff problems. The trade-off is that each implicit step requires solving a nonlinear system, which in turn requires the Jacobian — a matrix of partial derivatives of the RHS.

4.1 The Jacobian

The Jacobian \(J\) is an \(N \times N\) matrix where entry \(J_{i,j}\) describes how a small change in state component \(j\) affects the rate of change of component \(i\). The solver uses this sensitivity information to take stable, accurate steps through stiff regions. The components of the state vector \(\mathbf{u}\) — and therefore the size of \(J\) — are assembled in state_node_ids (core/src/util.jl):

one cumulative horizontal flow for each TabulatedRatingCurve, Pump, Outlet, LinearResistance, and ManningResistance node;
two cumulative horizontal flows per UserDemand (inflow and outflow tracked separately);
two cumulative vertical fluxes per Basin (evaporation and infiltration), which have their Basin as inflow node (inflow_id) and no outflow node, so they depend on the Basin storage through storage_uplink;
one integral term per PID control node.

Precipitation and drainage are deliberately not states — they are state-independent forcings and are integrated outside the ODE solver.

In practice, most entries of \(J\) are zero — a pump in one part of the network does not affect a weir in a distant part. This sparsity is the key to making implicit solves affordable, as both the Jacobian computation and the linear solves can exploit it.

The mathematics of how Ribasim evaluates the Jacobian and performs the linear solve is described in numerical considerations. This page describes the implementation, which lives in core/src/solve.jl.

5 Jacobian computation

The Jacobian computation pipeline lives in core/src/solve.jl.

5.1 Forward-mode automatic differentiation

There are three common ways to obtain a Jacobian: finite differences (simple but slow and imprecise), hand-coded derivatives (exact but hard to maintain), and automatic differentiation (AD). Ribasim uses forward-mode AD via ForwardDiff.jl, which propagates derivatives by replacing floating-point numbers with dual numbers — see the ForwardDiff documentation for the underlying theory. By running water_balance! with dual-number inputs seeded with a unit perturbation in one state direction, the derivative part of the output gives one column of the Jacobian; repeating for each direction yields the full matrix.

5.2 Exploiting sparsity: graph coloring

Computing each Jacobian column requires a full evaluation of water_balance! with dual numbers. With \(N\) states, that would mean \(N\) evaluations — prohibitively expensive for large models.

However, since the Jacobian is sparse, many columns have no overlapping nonzero rows. If columns \(i\) and \(j\) never have a nonzero in the same row, their perturbations can be seeded simultaneously in a single evaluation without interfering. This observation leads to graph coloring: columns are partitioned into groups (colors) such that all columns within a group are structurally independent. One evaluation per color suffices, reducing the cost from \(N\) evaluations to the number of colors — typically a small constant for network-structured problems.

5.3 Preparation phase

At model initialization, get_diff_eval prepares everything the solver will need. What is prepared depends on the solver setting reduced_implicit_solve.

With reduced_implicit_solve = true (default), RibasimJacobianEvaluationCache prepares two Jacobians:

∂flow_∂flow_input: the Jacobian of the flow formulation (formulate_flows_closure!) with respect to the input vector flow_input = (storage_uplink, storage_downlink, pid_integral, continuous_control_compound). Its output is the flow rates and the PID errors. This is \(\partial\mathbf{q}/\partial\mathbf{v}\) from the concept documentation;
∂continuous_control_compound_∂continuous_control_input: the Jacobian of the ContinuousControl compound variables with respect to (storage, flow).

With reduced_implicit_solve = false, the Jacobian of water_balance! with respect to the full state \(\mathbf{u}\) is prepared, and the result is a regular (sparse) matrix that OrdinaryDiffEq uses as is.

In both cases the preparation consists of:

Sparsity detection — TracerSparsityDetector from SparseConnectivityTracer runs the function with tracer inputs to determine which outputs depend on which inputs, producing a sparse boolean matrix.
Graph coloring — SparseMatrixColorings partitions the columns of the Jacobian into groups (colors) such that columns in the same group have no overlapping nonzero rows. This allows multiple columns to be computed in a single ForwardDiff pass by seeding multiple perturbation directions simultaneously. For ∂flow_∂flow_input the number of colors is small and independent of the model size, since its storage blocks are diagonal (except for flows controlled by PidControl with a derivative term).
AD preparation — DifferentiationInterface.prepare_jacobian pre-allocates all working arrays (dual-number buffers, etc.) based on the sparsity pattern and coloring.

backend_jac = AutoSparse(
    AutoForwardDiff(; chunksize, tag = :Ribasim);
    sparsity_detector = TracerSparsityDetector(),
)
∂flow_∂flow_input_prep = @ad_active p prepare_jacobian(
    formulate_flows_closure!,
    du,
    backend_jac,
    flow_input,
    Constant(t),
)

The closures capture the Parameters p. The @ad_active macro sets p.p_mutable.ad_active = true during AD calls, so that storage derived quantities are computed on demand instead of read from caches that would break the differentiable computational pipeline.

When sparse = false the AD backend is not wrapped in AutoSparse, and the Jacobians are dense matrices.

5.4 `Cache()` vs `Constant()` in DifferentiationInterface

When DifferentiationInterface.jacobian! calls the RHS, it needs to know how to handle each argument:

Cache(x) — “this depends on the differentiation variable; make a copy promoted to dual numbers.” The copy is allocated with the right dual-number element type but is not initialized with the original values. The RHS must write meaningful values into it before reading.
Constant(x) — “this is independent of the differentiation variable; pass the original object directly.” No copy, no promotion. But this means AD calls can mutate shared state.

This distinction is a performance optimization (fewer allocations), but it introduces a subtle invariant: any cached value in a Cache() argument must be freshly computed during the AD call, because the copy starts uninitialized. In the reduced_implicit_solve = false path and for the time derivative tgrad, storage_uplink, storage_downlink, the ContinuousControl compound variables and the current Basin storage are passed as Cache().

5.5 The `RibasimJacobian` operator

With reduced_implicit_solve = true, the Jacobian is never materialized. Instead the Jacobian prototype passed to the ODEFunction is a RibasimJacobian, a custom AbstractSciMLOperator which stores:

cache — the RibasimJacobianEvaluationCache with the AD preparation and the Jacobians described above;
J_inner_local — the \(n_\text{basin} \times n_\text{basin}\) matrix \(M Q_S\), the dependence of the Basin storage rates on the Basin storages via the flows (including via ContinuousControl);
area_pid_controlled — the area of the Basin listened to by each PidControl node.

The entries of \(Q_S\) (∂q_∂storage) are never stored, but iterated over by foreach_∂flow_∂storage, and similarly the dependence of the flows on the PID integrals by foreach_∂flow_∂pid_integral. These are used both to build the inner linear system and to expand its solution, which keeps the two consistent.

OrdinaryDiffEq calls update_coefficients! on the operator whenever the Jacobian may need updating. Whether the Jacobian should actually be re-evaluated is not passed to update_coefficients!, so Ribasim captures it from do_newJW in the p_mutable.refresh_jac flag. When refreshing, update_coefficients!:

Computes the current Basin storages and from those storage_uplink and storage_downlink (set_uplink_downlink_storage!);
Evaluates ∂flow_∂flow_input, which also caches the flows that ContinuousControl listens to;
Computes area_pid_controlled;
Evaluates ∂continuous_control_compound_∂continuous_control_input;
Computes J_inner_local (update_J_inner_local!).

6 Custom linear solver

At each Newton iteration the solver must solve a linear system involving the Jacobian. Ribasim provides a custom RibasimLinearSolve wrapper (in core/src/config.jl) around the actual factorization algorithm: KLUFactorization when sparse = true and LHLFactorization when sparse = false. With reduced_implicit_solve = false it falls back to the wrapped algorithm.

Otherwise SciMLBase.init sets up a RibasimLinearSolveCache, which contains an inner linear solve cache for the \(n_\text{basin} \times n_\text{basin}\) system in storage space, with an inner WOperator built from J_inner. The sparsity pattern of J_inner is determined once at initialization (with all contributions set to 1.0, so that entries where contributions cancel are kept) and kept constant, so the symbolic KLU factorization can be reused.

The linear solve itself is done in dolinsolve, following the concept documentation:

Reduce the right-hand side to storage space with aggregate_flows! (\(M\mathbf{b}^\text{flow}\)) plus the PID contribution;
Build J_inner with build_J_inner!: J_inner_local plus the PidControl contributions, which depend on \(\gamma\);
Solve the inner system \((J_\text{inner} - \gamma^{-1}I)\mathbf{s} = \mathbf{b}_\text{inner}\);
Expand the solution back to the full state space: the PID integral components follow from the change of the listened Basin storage, and the flow components from ∂flow_∂storage_mul! (\(Q_S\mathbf{s}\)).

This is significantly cheaper than solving in the full state space, since the inner system only has dimension equal to the number of Basins.

OrdinaryDiffEqCore.instability_jacobian returns J_inner for diagnostics when a simulation becomes unstable.

7 Callbacks

Ribasim uses SciML’s callback mechanism to handle events that occur during simulation. These are set up in core/src/callback.jl:

Callback	Type	Purpose
Negative storage check	`FunctionCallingCallback`	Detects and handles negative basin storage at every accepted step
Basin state saving	`SavingCallback`	Records basin storage/level at output times
Cumulative flow update	`FunctionCallingCallback`	Tracks cumulative flows for mass balance
Basin forcing update	`PresetTimeCallback`	Injects new forcing data at prescribed times
Flow saving	`SavingCallback`	Records averaged flows at output times
Discrete control	`FunctionCallingCallback`	Evaluates and applies discrete control rules
Tolerance decrease	`FunctionCallingCallback`	Compensates for growing cumulative state magnitudes (see numerical considerations)
Subgrid interpolation	`SavingCallback`	Interpolates subgrid levels for output

PresetTimeCallback fires at specific times (e.g. when forcing data changes). FunctionCallingCallback fires at every accepted solver step. SavingCallback fires at the saveat interval and stores results for later output.

1 SciML ecosystem overview

2 ODE problem setup

3 The right-hand side: water_balance!

3.1 Parameter structure

4 Implicit time integration

4.1 The Jacobian

5 Jacobian computation

5.1 Forward-mode automatic differentiation

5.2 Exploiting sparsity: graph coloring

5.3 Preparation phase

5.4 Cache() vs Constant() in DifferentiationInterface

5.5 The RibasimJacobian operator

6 Custom linear solver

7 Callbacks

3 The right-hand side: `water_balance!`

5.4 `Cache()` vs `Constant()` in DifferentiationInterface

5.5 The `RibasimJacobian` operator