Introduce DifferentiationInterface

Author: SouthEndMusic

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.11.3"
 manifest_format = "2.0"
-project_hash = "7591f8aefbf83678f96ece6aa98d0232f3c083ff"
+project_hash = "f1f6d1cb011ce63d93f72a0bbd984e2839b32e42"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "fb97701c117c8162e84dfcf80215caa904aef44f"
@@ -2083,7 +2083,7 @@ weakdeps = ["Distributed"]
     DistributedExt = "Distributed"
 
 [[deps.Ribasim]]
-deps = ["ADTypes", "Accessors", "Arrow", "BasicModelInterface", "CodecZstd", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "DiffEqBase", "DiffEqCallbacks", "EnumX", "FiniteDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "LineSearches", "LinearAlgebra", "LinearSolve", "Logging", "LoggingExtras", "MetaGraphsNext", "OrdinaryDiffEqBDF", "OrdinaryDiffEqCore", "OrdinaryDiffEqLowOrderRK", "OrdinaryDiffEqNonlinearSolve", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqTsit5", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "SparseConnectivityTracer", "StructArrays", "Tables", "TerminalLoggers", "TranscodingStreams"]
+deps = ["ADTypes", "Accessors", "Arrow", "BasicModelInterface", "CodecZstd", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "DiffEqBase", "DiffEqCallbacks", "DifferentiationInterface", "EnumX", "FiniteDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "LineSearches", "LinearAlgebra", "LinearSolve", "Logging", "LoggingExtras", "MetaGraphsNext", "OrdinaryDiffEqBDF", "OrdinaryDiffEqCore", "OrdinaryDiffEqLowOrderRK", "OrdinaryDiffEqNonlinearSolve", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqTsit5", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StructArrays", "Tables", "TerminalLoggers", "TranscodingStreams"]
 path = "core"
 uuid = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
 version = "2025.1.0"

Project.toml

@@ -18,6 +18,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 FindFirstFunctions = "64ca27bc-2ba2-4a57-88aa-44e436879224"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
@@ -57,6 +58,7 @@ SQLite = "0aa819cd-b072-5ff4-a722-6bc24af294d9"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"

core/Project.toml

@@ -17,6 +17,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
@@ -42,6 +43,7 @@ SQLite = "0aa819cd-b072-5ff4-a722-6bc24af294d9"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
@@ -63,6 +65,7 @@ DataStructures = "0.18"
 Dates = "1"
 DiffEqBase = "6.155"
 DiffEqCallbacks = "3.6, 4"
+DifferentiationInterface = "0.6.42"
 EnumX = "1.0"
 FiniteDiff = "2.21"
 Graphs = "1.9"
@@ -89,6 +92,7 @@ SQLite = "1.5.1"
 SciMLBase = "2.36"
 SparseArrays = "1"
 SparseConnectivityTracer = "0.6.8"
+SparseMatrixColorings = "0.4.13"
 Statistics = "1"
 StructArrays = "0.6.13, 0.7"
 TOML = "1"

core/src/Ribasim.jl

@@ -14,6 +14,10 @@ For more granular access, see:
 """
 module Ribasim
 
+using ADTypes: AutoForwardDiff, AutoFiniteDiff
+using DifferentiationInterface: AutoSparse, prepare_jacobian, jacobian!
+using SparseMatrixColorings: sparsity_pattern
+
 # Algorithms for solving ODEs.
 using OrdinaryDiffEqCore: OrdinaryDiffEqCore, get_du, AbstractNLSolver
 using DiffEqBase: DiffEqBase, calculate_residuals!
@@ -36,7 +40,7 @@ using SciMLBase:
 
 # Automatically detecting the sparsity pattern of the Jacobian of water_balance!
 # through operator overloading
-using SparseConnectivityTracer: TracerSparsityDetector, jacobian_sparsity, GradientTracer
+using SparseConnectivityTracer: GradientTracer
 
 # For efficient sparse computations
 using SparseArrays: SparseMatrixCSC, spzeros

core/src/config.jl

@@ -13,7 +13,6 @@ using DataStructures: DefaultDict
 using Dates: DateTime
 using Logging: LogLevel, Debug, Info, Warn, Error
 using ..Ribasim: Ribasim, isnode, nodetype
-using ADTypes: AutoForwardDiff, AutoFiniteDiff
 using OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, OrdinaryDiffEqNewtonAdaptiveAlgorithm
 using OrdinaryDiffEqNonlinearSolve: NLNewton
 using OrdinaryDiffEqLowOrderRK: Euler, RK4
@@ -310,10 +309,6 @@ function algorithm(solver::Solver; u0 = [])::OrdinaryDiffEqAlgorithm
         kwargs[:step_limiter!] = Ribasim.limit_flow!
     end
 
-    if function_accepts_kwarg(algotype, :autodiff)
-        kwargs[:autodiff] = solver.autodiff ? AutoForwardDiff() : AutoFiniteDiff()
-    end
-
     algotype(; kwargs...)
 end
 

core/src/model.jl

@@ -28,6 +28,33 @@ struct Model{T}
     end
 end
 
+function get_jac_eval(du::Vector, u::Vector, p::Parameters, solver::Solver)
+    backend = if solver.autodiff
+        AutoForwardDiff()
+    else
+        AutoFiniteDiff()
+    end
+
+    if solver.sparse
+        backend = AutoSparse(backend)
+    end
+
+    t = 0.0
+
+    # Activate all nodes to catch all possible state dependencies
+    p.all_nodes_active = true
+    prep = prepare_jacobian((du, u) -> water_balance(du, u, p, t), du, backend, u)
+    p.all_nodes_active = false
+
+    jac_prototype = sparsity_pattern(prep)
+
+    jac =
+        (J, u, p, t) ->
+            jacobian!((du, u) -> water_balance!(du, u, p, t), du, J, prep, backend, u)
+
+    return jac_prototype, jac
+end
+
 function Model(config_path::AbstractString)::Model
     config = Config(config_path)
     if !valid_config(config)
@@ -159,17 +186,13 @@ function Model(config::Config)::Model
     tstops = sort(unique(vcat(tstops...)))
     adaptive, dt = convert_dt(config.solver.dt)
 
-    jac_prototype = if config.solver.sparse
-        get_jac_prototype(du0, u0, parameters, t0)
-    else
-        nothing
-    end
-    RHS = ODEFunction(water_balance!; jac_prototype)
+    jac_prototype, jac = get_jac_eval(du0, u0, parameters, config.solver)
+    RHS = ODEFunction(water_balance!; jac_prototype, jac)
 
     prob = ODEProblem(RHS, u0, timespan, parameters)
     @debug "Setup ODEProblem."
 
-    callback, saved = create_callbacks(parameters, config, u0, saveat)
+    callback, saved = create_callbacks(parameters, config, saveat)
     @debug "Created callbacks."
 
     # Run water_balance! before initializing the integrator. This is because

core/src/util.jl

@@ -764,18 +764,6 @@ function relaxed_root(x, threshold)
     end
 end
 
-function get_jac_prototype(du0, u0, p, t0)
-    p.all_nodes_active = true
-    jac_prototype = jacobian_sparsity(
-        (du, u) -> water_balance!(du, u, p, t0),
-        du0,
-        u0,
-        TracerSparsityDetector(),
-    )
-    p.all_nodes_active = false
-    jac_prototype
-end
-
 # Custom overloads
 reduction_factor(x::GradientTracer, ::Real) = x
 low_storage_factor_resistance_node(storage, q::GradientTracer, inflow_id, outflow_id) = q