Use mixed mode AD for Jac evaluation

Project.toml

@@ -43,6 +43,7 @@ LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 MarkdownTables = "1862ce21-31c7-451e-824c-f20fa3f90fa2"
 MathOptAnalyzer = "d1179b25-476b-425c-b826-c7787f0fff83"
 MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Moshi = "2e0e35c7-a2e4-4343-998d-7ef72827ed2d"
 MultiObjectiveAlgorithms = "0327d340-17cd-11ea-3e99-2fd5d98cecda"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"

core/Project.toml

@@ -31,6 +31,7 @@ Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 MathOptAnalyzer = "d1179b25-476b-425c-b826-c7787f0fff83"
 MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Moshi = "2e0e35c7-a2e4-4343-998d-7ef72827ed2d"
 MultiObjectiveAlgorithms = "0327d340-17cd-11ea-3e99-2fd5d98cecda"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
@@ -86,6 +87,7 @@ Logging = "1"
 LoggingExtras = "1"
 MathOptAnalyzer = "0.1.0"
 MetaGraphsNext = "0.6, 0.7"
+Mooncake = "0.4.168"
 Moshi = "0.3.7"
 MultiObjectiveAlgorithms = "1.6.0"
 NCDatasets = "0.14.8"

core/src/Ribasim.jl

@@ -19,19 +19,14 @@ using Preferences: @load_preference
 
 # Requirements for automatic differentiation
 using DifferentiationInterface:
-    AutoSparse,
-    Constant,
-    Cache,
-    prepare_jacobian,
-    jacobian!,
-    prepare_derivative,
-    derivative!
+    Constant, Cache, prepare_jacobian, jacobian!, prepare_derivative, derivative!
 
 using ForwardDiff: derivative as forward_diff
 
 # Algorithms for solving ODEs.
 using OrdinaryDiffEqCore: OrdinaryDiffEqCore, get_du
 import ForwardDiff
+import Mooncake
 
 # Interface for defining and solving the ODE problem of the physical layer.
 using SciMLBase:
@@ -49,8 +44,15 @@ using SciMLBase:
 
 # Automatically detecting the sparsity pattern of the Jacobian of water_balance!
 # through operator overloading
-using SparseConnectivityTracer: GradientTracer, TracerSparsityDetector
-using SparseMatrixColorings: GreedyColoringAlgorithm, sparsity_pattern
+using SparseConnectivityTracer: GradientTracer
+using SparseMatrixColorings:
+    NaturalOrder,
+    LargestFirst,
+    SmallestLast,
+    IncidenceDegree,
+    DynamicLargestFirst,
+    DynamicDegreeBasedOrder,
+    sparsity_pattern
 
 # For efficient sparse computations
 using SparseArrays: SparseMatrixCSC, spzeros, sparse

core/src/allocation_optim.jl

@@ -16,7 +16,7 @@ function set_simulation_data!(
         user_demand,
         tabulated_rating_curve,
     ) = p.p_independent
-    du = get_du(integrator)
+    du = wrap_state(get_du(integrator), p.p_independent)
 
     errors = false
 
@@ -577,7 +577,7 @@ function warm_start!(allocation_model::AllocationModel, integrator::DEIntegrator
     (; basin_ids_subnetwork) = node_ids_in_subnetwork
     flow = problem[:flow]
     storage_change = problem[:basin_storage_change]
-    du = get_du(integrator)
+    du = wrap_state(get_du(integrator), p.p_independent)
 
     # Extrapolate the current instantaneous flow rates from the physical layer
     for link in only(flow.axes)

core/src/bmi.jl

@@ -57,9 +57,9 @@ end
 This uses a typeassert to ensure that the return type annotation doesn't create a copy.
 """
 function BMI.get_value_ptr(model::Model, name::String)::Vector{Float64}
-    (; u, p) = model.integrator
-    (; p_independent, state_time_dependent_cache) = p
+    (; p_independent, state_time_dependent_cache) = model.integrator.p
     (; basin, user_demand, subgrid) = p_independent
+    u = Ribasim.get_wrapped_u(model)
 
     if name == "basin.storage"
         state_time_dependent_cache.current_storage

core/src/callback.jl

@@ -86,8 +86,9 @@ end
 Decrease the relative tolerance of the integrator over time,
 to compensate for the ever increasing cumulative flows.
 """
-function decrease_tolerance!(u, t, integrator)::Nothing
+function decrease_tolerance!(u_raw, t, integrator)::Nothing
     (; p, t, opts) = integrator
+    u = wrap_state(u_raw, p.p_independent)
 
     for (i, state) in enumerate(u)
         p.p_independent.relmask[i] || continue
@@ -121,11 +122,12 @@ as each flow carries mass, based on the concentrations of the flow source.
 Specifically, we first use all the inflows to update the mass of the Basins, recalculate
 the Basin concentration(s) and then remove the mass that is being lost to the outflows.
 """
-function update_cumulative_flows!(u, t, integrator)::Nothing
+function update_cumulative_flows!(u_raw, t, integrator)::Nothing
     (; cache, p) = integrator
     (; p_independent, p_mutable, time_dependent_cache) = p
     (; basin, flow_boundary, allocation, temp_convergence, convergence, ncalls) =
         p_independent
+    u = wrap_state(u_raw, p_independent)
 
     # Update tprev
     p_mutable.tprev = t
@@ -196,8 +198,8 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
     return nothing
 end
 
-function update_concentrations!(u, t, integrator)::Nothing
-    (; uprev, p, tprev, dt) = integrator
+function update_concentrations!(u_raw, t, integrator)::Nothing
+    (; p, tprev, dt) = integrator
     (; p_independent, state_time_dependent_cache) = p
     (; current_storage, current_level) = state_time_dependent_cache
     (; basin, flow_boundary, do_concentration) = p_independent
@@ -211,6 +213,8 @@ function update_concentrations!(u, t, integrator)::Nothing
         concentration_itp_surface_runoff,
         mass,
     ) = concentration_data
+    u = wrap_state(u_raw, p_independent)
+    uprev = wrap_state(integrator.uprev, p_independent)
 
     !do_concentration && return nothing
 
@@ -336,10 +340,12 @@ end
 Compute the forcing volume entering and leaving the Basin over the last time step
 """
 function forcing_update(integrator::DEIntegrator, node_id::NodeID)::Tuple{Float64, Float64}
-    (; u, uprev, p, dt) = integrator
+    (; p, dt) = integrator
     (; basin) = p.p_independent
     (; vertical_flux) = basin
 
+    u = wrap_state(integrator.u, p.p_independent)
+    uprev = wrap_state(integrator.uprev, p.p_independent)
     @assert node_id.type == NodeType.Basin
 
     fixed_area = basin_areas(basin, node_id.idx)[end]
@@ -366,8 +372,10 @@ function flow_update_on_link(
     integrator::DEIntegrator,
     link_src::Tuple{NodeID, NodeID},
 )::Float64
-    (; u, uprev, p, t, tprev) = integrator
+    (; p, t, tprev) = integrator
     (; flow_boundary) = p.p_independent
+    u = wrap_state(integrator.u, p.p_independent)
+    uprev = wrap_state(integrator.uprev, p.p_independent)
 
     from_id, to_id = link_src
     if from_id == to_id
@@ -390,7 +398,7 @@ end
 """
 Save the storages and levels at the latest t.
 """
-function save_basin_state(u, t, integrator)
+function save_basin_state(u_raw, t, integrator)
     (; current_storage, current_level) = integrator.p.state_time_dependent_cache
     SavedBasinState(; storage = copy(current_storage), level = copy(current_level), t)
 end
@@ -400,7 +408,7 @@ Save all cumulative forcings and flows over links over the latest timestep,
 Both computed by the solver and integrated exactly. Also computes the total horizontal
 inflow and outflow per Basin.
 """
-function save_flow(u, t, integrator)
+function save_flow(u_raw, t, integrator)
     (; cache, p) = integrator
     (;
         basin,
@@ -412,6 +420,7 @@ function save_flow(u, t, integrator)
         ncalls,
         node_id,
     ) = p.p_independent
+    u = wrap_state(u_raw, p.p_independent)
     Δt = get_Δt(integrator)
     flow_mean = (u - u_prev_saveat) / Δt
 
@@ -504,12 +513,13 @@ function check_water_balance_error!(
     integrator::DEIntegrator,
     Δt::Float64,
 )::Nothing
-    (; u, p, t) = integrator
+    (; p, t) = integrator
     (; p_independent, state_time_dependent_cache) = p
     (; current_storage) = state_time_dependent_cache
-    (; basin, water_balance_abstol, water_balance_reltol, starttime) = p_independent
+    (; basin, water_balance_abstol, water_balance_reltol, starttime, state_ranges) =
+        p_independent
     errors = false
-    state_ranges = getaxes(u)
+    u = wrap_state(integrator.u, p_independent)
 
     # The initial storage is irrelevant for the storage rate and can only cause
     # floating point truncation errors
@@ -568,7 +578,7 @@ function check_water_balance_error!(
     return nothing
 end
 
-function save_solver_stats(u, t, integrator)
+function save_solver_stats(u_raw, t, integrator)
     (; dt) = integrator
     (; stats) = integrator.sol
     (;
@@ -582,12 +592,11 @@ function save_solver_stats(u, t, integrator)
     )
 end
 
-function check_negative_storage(u, t, integrator)::Nothing
+function check_negative_storage(u_raw, t, integrator)::Nothing
     (; p) = integrator
     (; p_independent, state_time_dependent_cache) = p
     (; basin) = p_independent
-    du = get_du(integrator)
-    water_balance!(du, u, p, t)
+    set_current_basin_properties!(wrap_state(u_raw, p_independent), p, t)
 
     errors = false
     for id in basin.node_id
@@ -616,11 +625,11 @@ Apply the discrete control logic. There's somewhat of a complex structure:
 - The nodes that are controlled by this DiscreteControl node must have the same control state, for which they have
     parameter values associated with that control state defined in their control_mapping
 """
-function apply_discrete_control!(u, t, integrator)::Nothing
+function apply_discrete_control!(u_raw, t, integrator)::Nothing
     (; p) = integrator
     (; discrete_control) = p.p_independent
     (; node_id, truth_state, compound_variables) = discrete_control
-    du = get_du(integrator)
+    du = wrap_state(get_du(integrator), p.p_independent)
 
     # Loop over the discrete control nodes to determine their truth state
     # and detect possible control state changes
@@ -723,7 +732,7 @@ function get_value(subvariable::SubVariable, p::Parameters, du::CVector, t::Floa
             level = level_boundary.level[listen_node_id.idx](t + look_ahead)
         else
             error(
-                "Level condition node '$node_id' is neither a basin nor a level boundary.",
+                "Level condition node '$listen_node_id' is neither a basin nor a level boundary.",
             )
         end
         value = level

core/src/carrays.jl

@@ -45,7 +45,7 @@ Base.size(x::CArray) = size(getdata(x))
 Base.length(x::CArray) = length(getdata(x))
 Base.getindex(x::CArray, i::Int) = getdata(x)[i]
 Base.getindex(x::CArray, I...) = getdata(x)[I...]
-Base.IndexStyle(::Type{CArray}) = IndexLinear()
+Base.IndexStyle(::Type{<:CArray}) = IndexLinear()
 Base.elsize(x::CArray) = Base.elsize(getdata(x))
 
 # Linear algebra

core/src/config.jl

@@ -8,10 +8,11 @@ Ribasim.config is a submodule mainly to avoid name clashes between the configura
 """
 module config
 
-using ADTypes: AutoForwardDiff, AutoFiniteDiff
+using ADTypes: AutoForwardDiff, AutoMooncake, AutoFiniteDiff
 using Configurations: Configurations, @option, from_toml, @type_alias
 using DataStructures: OrderedDict
 using Dates: DateTime
+using DifferentiationInterface: AutoSparse, MixedMode
 using Logging: LogLevel, Debug, Info, Warn, Error
 using ..Ribasim: Ribasim, Table, Schema
 using OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, OrdinaryDiffEqNewtonAdaptiveAlgorithm
@@ -21,12 +22,15 @@ using OrdinaryDiffEqTsit5: Tsit5
 using OrdinaryDiffEqSDIRK: ImplicitEuler, KenCarp4, TRBDF2
 using OrdinaryDiffEqBDF: FBDF, QNDF
 using OrdinaryDiffEqRosenbrock: Rosenbrock23, Rodas4P, Rodas5P
+using SparseConnectivityTracer: TracerSparsityDetector
+using SparseMatrixColorings: GreedyColoringAlgorithm, AbstractOrder, NaturalOrder
 using LinearSolve: KLUFactorization
 
 export Config, Solver, Results, Logging, Toml
 export algorithm,
     camel_case,
-    get_ad_type,
+    get_jac_ad_backend,
+    get_tgrad_ad_backend,
     snake_case,
     input_path,
     database_path,
@@ -171,6 +175,7 @@ end
 @option struct Experimental <: TableOption
     concentration::Bool = false
     allocation::Bool = false
+    optimize_jacobian::Bool = false
 end
 
 # For logging enabled experimental features
@@ -358,16 +363,48 @@ function function_accepts_kwarg(f, kwarg)::Bool
     return false
 end
 
-function get_ad_type(solver::Solver; specialize = true)
+get_forward_mode_ad_type(; specialize::Bool = true) =
+    AutoForwardDiff(; chunksize = specialize ? nothing : 1, tag = :Ribasim_forward)
+
+function get_jac_ad_backend(
+    solver::Solver;
+    mixed_mode::Bool = false,
+    specialize = true,
+    order::Union{Function, Type{<:AbstractOrder}} = NaturalOrder,
+)
+    forward_mode =
+        solver.autodiff ? get_forward_mode_ad_type(; specialize) : AutoFiniteDiff()
+    if solver.sparse
+        sparsity_detector = TracerSparsityDetector()
+        backend = if mixed_mode
+            reverse_mode = AutoMooncake()
+            backend = MixedMode(forward_mode, reverse_mode)
+        else
+            forward_mode
+        end
+        AutoSparse(
+            backend;
+            sparsity_detector,
+            coloring_algorithm = GreedyColoringAlgorithm{:substitution}(
+                order();
+                postprocessing = true,
+            ),
+        )
+    else
+        forward_mode
+    end
+end
+
+function get_tgrad_ad_backend(solver::Solver; specialize = true)
     if solver.autodiff
-        AutoForwardDiff(; chunksize = specialize ? nothing : 1, tag = :Ribasim)
+        get_forward_mode_ad_type(; specialize)
     else
         AutoFiniteDiff()
     end
 end
 
 "Create an OrdinaryDiffEqAlgorithm from solver config"
-function algorithm(solver::Solver; u0 = [], specialize = true)::OrdinaryDiffEqAlgorithm
+function algorithm(solver::Solver; specialize = true)::OrdinaryDiffEqAlgorithm
     kwargs = Dict{Symbol, Any}()
     algotype = algorithms[solver.algorithm]
 
@@ -382,10 +419,6 @@ function algorithm(solver::Solver; u0 = [], specialize = true)::OrdinaryDiffEqAl
         kwargs[:step_limiter!] = Ribasim.limit_flow!
     end
 
-    if function_accepts_kwarg(algotype, :autodiff)
-        kwargs[:autodiff] = get_ad_type(solver; specialize)
-    end
-
     algotype(; kwargs...)
 end