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.

Which of these the solver actually uses depends on the algorithm:

Implicit algorithms with a Newton solver (NordsieckBDF, FBDF, QNDF, TRBDF2, KenCarp4, ImplicitEuler) need the Jacobian in their nonlinear solve;
Rosenbrock algorithms (Rosenbrock23, Rodas4P, Rodas5P) need the full Jacobian and additionally the time derivative tgrad of the right hand side;
Explicit algorithms (Tsit5, RK4, Euler) need neither. get_diff_eval is currently called for every algorithm, so the preparation below is also done for them, but the result is not used.

The Jacobian is prepared in one of two forms, set by reduced_implicit_solve. This is true by default for the implicit algorithms with a Newton solver, and always false for Rosenbrock and explicit algorithms.

With reduced_implicit_solve = true, RibasimJacobianEvaluationCache prepares three 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);
∂junction_flow_∂flow: the Jacobian \(K\) of formulate_junction_flow!, the last step of water_balance!, which computes the flows of the UserDemand nodes that abstract from Junction nodes from the other flows. It is a square matrix over all flows, with an identity part for all other flows.

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. The time derivative tgrad is prepared in the same way, for the Rosenbrock algorithms.

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(). For tgrad the TimeDependentCache is also passed as Cache(), since the time t is the differentiation variable and the cached time dependent values must be dual numbers too. Therefore the components of Parameters are separate arguments of the wrapped right hand side, which packs them back into a Parameters with the dual TimeDependentCache.

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 K Q_S\), the dependence of the Basin storage rates on the Basin storages via the flows (including via ContinuousControl and the routing through Junction subgraphs, which is in \(M\));
area_pid_controlled — the area of the Basin listened to by each PidControl node;
flow_buffer — a work vector for the linear solve.

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. The contributions of the flow in these entries to the Basin storages are added for every flow that depends on it via \(K\), see foreach_∂flow_∂flow_np. Without Junction nodes that stay in the simulation \(K\) is the identity and this is skipped (has_routed_junctions). 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;
Evaluates ∂junction_flow_∂flow at the flows cached in step 2, if the model has Junction nodes that stay in the simulation;
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, which is passed through \(K\) if the model has Junction nodes that stay in the simulation;
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}\)), where the flows in Junction subgraphs follow by applying \(K\) (∂flow_∂flow_np_mul!).

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 Junction nodes that stay in the simulation

A Junction with rows in Junction / static is “routed” in the code (NodeID.is_routed, Junction.routed_ids): it divides the flow or supplies a UserDemand, and stays in the graph, while the other Junctions are removed by simplify_graph!. Whether a Junction is routed is decided when the node ids are read (routed_junction_ids), so the node types and validation can tell the two apart before the graph is simplified. A routed Junction may flow into a plain Junction: replaced_link_targets follows flow_link_map to the link that replaced the drawn link, so the row applies to the node behind the plain Junction. If the plain Junction divides the flow over several nodes, the row is an error.

The links in the subgraphs of routed Junctions have no states. set_routing! computes at initialization, per link out of a routed Junction and per routed Junction, the weights of the flows that enter the subgraph (the states of connector nodes and UserDemand outflows, and the FlowBoundary flows) and of the inflows of the UserDemand nodes in the subgraph (RoutingWeights), and per flow state and per FlowBoundary the Basins they reach through the subgraphs with their weights (routed_basins, routed_basins_boundary). These weights enter the incidence matrix, aggregate_flows!, the storage reconstruction and the Jacobian (add_flow_contribution!), and routed_flow evaluates the flow over a routed link for the results, the water balance check and the tracers. Only the UserDemand flows in the subgraphs are computed by formulate_junction_flow!, the last step of water_balance!, see the concept documentation. Two consequences are worth knowing.

The downstream level is looked up per flow. downlink_level_id holds per flow state the node whose level the structure sees downstream (set_downstream_level_ids!): the node downstream, or behind a chain of routed Junctions that do not divide the flow the single Basin or LevelBoundary at its end (Junction.level_node_id). It fills storage_downlink and places the downlink derivatives in foreach_∂flow_∂storage, while outflow_id and the incidence matrix keep the Junction, since the volume reaches the Basins through the routing weights.

Boundary flows stay exact. The flow of a FlowBoundary into a Basin is never a state: its cumulative volume is integrated exactly and added to the storage via the time dependent cache, so the solver never sees the time series. With the routing this also holds for a FlowBoundary into a routed Junction: the Basins behind the subgraph get their share of the exact cumulative flow. The UserDemand nodes in the subgraph do see the current rate, which set_current_boundary_flow_rate! caches like the other time series in the right hand side: a flow rate set via BMI takes precedence, and a block series is evaluated left-continuously within the timestep. A daily on/off FlowBoundary over a month takes 37 steps through a routed Junction, the same as straight into a Basin. When the links in the subgraphs were states, it took about 130 accepted and 90 rejected steps, and thousands when the series was evaluated at t directly.

PID control. The derivative term of PID control uses the storage rate of the listened Basin, which includes the flows into it via routed Junctions. Since those are computed after the controlled flows, formulate_provisional_junction_flow! computes them once before PID control from the flows known at that point, in water_balance! and in the Jacobian closure alike. The flows of PID and ContinuousControl controlled nodes are still zero then, so a controlled flow that reaches the listened Basin via routed Junctions is not part of the derivative term, in the same way that a controlled flow is only accounted for in the derivative term when it flows into the listened Basin directly.

ContinuousControl. ContinuousControl can listen to the flow rate of a routed Junction or of a UserDemand that abstracts from one, which depend on the UserDemand flows in the subgraph. formulate_provisional_junction_flow! therefore computes those once more before the compound variables, after the PID controlled flows. The flows of ContinuousControl controlled nodes are still zero then, as for any flow ContinuousControl listens to. In the Jacobian the closure of the compound variables computes the flows in the subgraphs from its input flows first, so the compound variables are differentiated with respect to the flows that enter the subgraphs, which is where foreach_∂flow_∂storage picks up their dependence on the storages.

8 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
Cumulative flow update	`FunctionCallingCallback`	Tracks cumulative flows for mass balance
Basin forcing update	`PresetTimeCallback`	Injects new forcing data at prescribed times
Result saving	`SavingCallback`	Records all results over time at output times
Discrete control	`FunctionCallingCallback`	Evaluates and applies discrete control rules
Tolerance decrease	`FunctionCallingCallback`	Compensates for growing cumulative state magnitudes (see numerical considerations)

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 Junction nodes that stay in the simulation

8 Callbacks

3 The right-hand side: `water_balance!`

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

5.5 The `RibasimJacobian` operator