Introduce differentiation interface

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.11.3"
 manifest_format = "2.0"
-project_hash = "7591f8aefbf83678f96ece6aa98d0232f3c083ff"
+project_hash = "08df9cf37f934d9c54c355f731e7bb06c158c106"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "fb97701c117c8162e84dfcf80215caa904aef44f"
@@ -547,9 +547,9 @@ version = "1.15.1"
 
 [[deps.DifferentiationInterface]]
 deps = ["ADTypes", "LinearAlgebra"]
-git-tree-sha1 = "258fa016b2d03f19e4d0d1cd8e30c84907af1528"
+git-tree-sha1 = "cf6dcb4b21bdd893fd45be70791d6dc89ca16506"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
-version = "0.6.42"
+version = "0.6.48"
 
     [deps.DifferentiationInterface.extensions]
     DifferentiationInterfaceChainRulesCoreExt = "ChainRulesCore"
@@ -561,9 +561,10 @@ version = "0.6.42"
     DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"]
     DifferentiationInterfaceGTPSAExt = "GTPSA"
     DifferentiationInterfaceMooncakeExt = "Mooncake"
-    DifferentiationInterfacePolyesterForwardDiffExt = "PolyesterForwardDiff"
+    DifferentiationInterfacePolyesterForwardDiffExt = ["PolyesterForwardDiff", "ForwardDiff", "DiffResults"]
     DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
     DifferentiationInterfaceSparseArraysExt = "SparseArrays"
+    DifferentiationInterfaceSparseConnectivityTracerExt = "SparseConnectivityTracer"
     DifferentiationInterfaceSparseMatrixColoringsExt = "SparseMatrixColorings"
     DifferentiationInterfaceStaticArraysExt = "StaticArrays"
     DifferentiationInterfaceSymbolicsExt = "Symbolics"
@@ -585,6 +586,7 @@ version = "0.6.42"
     PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
     ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
     SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+    SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
     SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
     StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
     Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
@@ -2083,7 +2085,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", "SQLite", "SciMLBase", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StructArrays", "Tables", "TerminalLoggers", "TranscodingStreams"]
 path = "core"
 uuid = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
 version = "2025.2.0"
@@ -2291,9 +2293,9 @@ version = "1.11.0"
 
 [[deps.SparseConnectivityTracer]]
 deps = ["ADTypes", "DocStringExtensions", "FillArrays", "LinearAlgebra", "Random", "SparseArrays"]
-git-tree-sha1 = "6651f4663027f3b30a31429d257185f56a571184"
+git-tree-sha1 = "e60ba2f27bf0c85cc20ac326a158182630032bb6"
 uuid = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
-version = "0.6.13"
+version = "0.6.14"
 
     [deps.SparseConnectivityTracer.extensions]
     SparseConnectivityTracerDataInterpolationsExt = "DataInterpolations"
@@ -2331,9 +2333,9 @@ version = "2.23.1"
 
 [[deps.SparseMatrixColorings]]
 deps = ["ADTypes", "DataStructures", "DocStringExtensions", "LinearAlgebra", "Random", "SparseArrays"]
-git-tree-sha1 = "97092c0a40d6033b7da27ea15bcf75fd5b446254"
+git-tree-sha1 = "e0ae9189392572abe85bc9fd4ce35e772b1e1e10"
 uuid = "0a514795-09f3-496d-8182-132a7b665d35"
-version = "0.4.13"
+version = "0.4.14"
 weakdeps = ["Colors"]
 
     [deps.SparseMatrixColorings.extensions]

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"
@@ -50,13 +51,13 @@ OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 OteraEngine = "b2d7f28f-acd6-4007-8b26-bc27716e5513"
 PackageCompiler = "9b87118b-4619-50d2-8e1e-99f35a4d4d9d"
-PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 Ribasim = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
 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"
@@ -37,11 +38,11 @@ OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
 OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
 OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
-PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 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 +64,7 @@ DataStructures = "0.18"
 Dates = "1"
 DiffEqBase = "6.155"
 DiffEqCallbacks = "3.6, 4"
+DifferentiationInterface = "0.6.48"
 EnumX = "1.0"
 FiniteDiff = "2.21"
 Graphs = "1.9"
@@ -84,11 +86,11 @@ OrdinaryDiffEqNonlinearSolve = "1.1"
 OrdinaryDiffEqRosenbrock = "1.1"
 OrdinaryDiffEqSDIRK = "1.1"
 OrdinaryDiffEqTsit5 = "1.1"
-PreallocationTools = "0.4"
 SQLite = "1.5.1"
 SciMLBase = "2.36"
 SparseArrays = "1"
 SparseConnectivityTracer = "0.6.8"
+SparseMatrixColorings = "0.4.14"
 Statistics = "1"
 StructArrays = "0.6.13, 0.7"
 TOML = "1"

core/src/Ribasim.jl

@@ -14,6 +14,16 @@ For more granular access, see:
 """
 module Ribasim
 
+# Requirements for automatic differentiation
+using DifferentiationInterface:
+    AutoSparse,
+    Constant,
+    Cache,
+    prepare_jacobian,
+    jacobian!,
+    prepare_derivative,
+    derivative!
+
 # Algorithms for solving ODEs.
 using OrdinaryDiffEqCore: OrdinaryDiffEqCore, get_du, AbstractNLSolver
 using DiffEqBase: DiffEqBase, calculate_residuals!
@@ -36,23 +46,16 @@ 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, TracerSparsityDetector, IndexSetGradientPattern
+using SparseMatrixColorings: GreedyColoringAlgorithm, sparsity_pattern
 
 # For efficient sparse computations
 using SparseArrays: SparseMatrixCSC, spzeros
 
 # Linear algebra
 using LinearAlgebra: mul!
 
-# PreallocationTools is used because the RHS function (water_balance!) gets called with different input types
-# for u, du:
-# - Float64 for normal calls
-# - Dual numbers for automatic differentiation with ForwardDiff
-# - GradientTracer for automatic Jacobian sparsity detection with SparseConnectivityTracer
-# The computations inside the rhs go trough preallocated arrays of the required type which are created by LazyBufferCache.
-# Retrieving a cache from a LazyBufferCache looks like indexing: https://docs.sciml.ai/PreallocationTools/stable/#LazyBufferCache
-using PreallocationTools: LazyBufferCache
-
 # Interpolation functionality, used for e.g.
 # basin profiles and TabulatedRatingCurve. See also the node
 # references in the docs.
@@ -65,7 +68,8 @@ using DataInterpolations:
     integral,
     AbstractInterpolation,
     ExtrapolationType
-using DataInterpolations.ExtrapolationType: Constant, Periodic, Extension, Linear
+using DataInterpolations.ExtrapolationType:
+    Constant as ConstantExtrapolation, Periodic, Extension, Linear
 
 # Modeling language for Mathematical Optimization.
 # Used for allocation, see the docs: https://ribasim.org/dev/allocation.html

core/src/allocation_init.jl

@@ -4,10 +4,10 @@ which is a type of sparse arrays that in this case takes NodeID in stead of Int
 E.g. capacity[(node_a, node_b)] gives the capacity of link (node_a, node_b).
 """
 function get_subnetwork_capacity(
-    p::Parameters,
+    p_non_diff::ParametersNonDiff,
     subnetwork_id::Int32,
 )::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}}
-    (; graph) = p
+    (; graph) = p_non_diff
     node_ids_subnetwork = graph[].node_ids[subnetwork_id]
 
     dict = Dict{Tuple{NodeID, NodeID}, Float64}()
@@ -28,7 +28,7 @@ function get_subnetwork_capacity(
 
             # Find flow constraints for this link
             if is_flow_constraining(id_src.type)
-                node_src = getfield(p, graph[id_src].type)
+                node_src = getfield(p_non_diff, graph[id_src].type)
 
                 max_flow_rate = node_src.max_flow_rate[id_src.idx]
                 capacity_node_src =
@@ -37,7 +37,7 @@ function get_subnetwork_capacity(
                 capacity_link = min(capacity_link, capacity_node_src)
             end
             if is_flow_constraining(id_dst.type)
-                node_dst = getfield(p, graph[id_dst].type)
+                node_dst = getfield(p_non_diff, graph[id_dst].type)
 
                 max_flow_rate = node_dst.max_flow_rate[id_dst.idx]
                 capacity_node_dst =
@@ -72,10 +72,10 @@ and subnetwork allocation network (defined by their capacity objects).
 """
 function add_subnetwork_connections!(
     capacity::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}},
-    p::Parameters,
+    p_non_diff::ParametersNonDiff,
     subnetwork_id::Int32,
 )::Nothing
-    (; allocation) = p
+    (; allocation) = p_non_diff
     (; main_network_connections) = allocation
 
     # Add the connections to the main network
@@ -100,11 +100,11 @@ dictionary wrapper. The keys of this dictionary define
 the which links are used in the allocation optimization problem.
 """
 function get_capacity(
-    p::Parameters,
+    p_non_diff::ParametersNonDiff,
     subnetwork_id::Int32,
 )::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}}
-    capacity = get_subnetwork_capacity(p, subnetwork_id)
-    add_subnetwork_connections!(capacity, p, subnetwork_id)
+    capacity = get_subnetwork_capacity(p_non_diff, subnetwork_id)
+    add_subnetwork_connections!(capacity, p_non_diff, subnetwork_id)
 
     return capacity
 end
@@ -176,10 +176,10 @@ flow over link <= link capacity
 function add_constraints_capacity!(
     problem::JuMP.Model,
     capacity::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}},
-    p::Parameters,
+    p_non_diff::ParametersNonDiff,
     subnetwork_id::Int32,
 )::Nothing
-    (; main_network_connections) = p.allocation
+    (; main_network_connections) = p_non_diff.allocation
     main_network_source_links = main_network_connections[subnetwork_id]
     F = problem[:F]
 
@@ -291,10 +291,10 @@ sum(flows out of basin) == sum(flows into basin) + flow from storage and vertica
 """
 function add_constraints_conservation_node!(
     problem::JuMP.Model,
-    p::Parameters,
+    p_non_diff::ParametersNonDiff,
     subnetwork_id::Int32,
 )::Nothing
-    (; graph) = p
+    (; graph) = p_non_diff
     F = problem[:F]
     F_basin_in = problem[:F_basin_in]
     F_basin_out = problem[:F_basin_out]
@@ -405,7 +405,7 @@ end
 Construct the allocation problem for the current subnetwork as a JuMP model.
 """
 function allocation_problem(
-    p::Parameters,
+    p_non_diff::ParametersNonDiff,
     sources::OrderedDict{Tuple{NodeID, NodeID}, AllocationSource},
     capacity::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}},
     subnetwork_id::Int32,
@@ -427,8 +427,8 @@ function allocation_problem(
     add_variables_flow_buffer!(problem, sources)
 
     # Add constraints to problem
-    add_constraints_conservation_node!(problem, p, subnetwork_id)
-    add_constraints_capacity!(problem, capacity, p, subnetwork_id)
+    add_constraints_conservation_node!(problem, p_non_diff, subnetwork_id)
+    add_constraints_capacity!(problem, capacity, p_non_diff, subnetwork_id)
     add_constraints_boundary_source!(problem, sources)
     add_constraints_main_network_source!(problem, sources)
     add_constraints_user_source!(problem, sources)
@@ -443,17 +443,17 @@ Construct the JuMP.jl problem for allocation.
 """
 function AllocationModel(
     subnetwork_id::Int32,
-    p::Parameters,
+    p_non_diff::ParametersNonDiff,
     sources::OrderedDict{Tuple{NodeID, NodeID}, AllocationSource},
     Δt_allocation::Float64,
 )::AllocationModel
-    capacity = get_capacity(p, subnetwork_id)
+    capacity = get_capacity(p_non_diff, subnetwork_id)
     source_priorities = unique(source.source_priority for source in values(sources))
     sources = OrderedDict(
         link => source for
         (link, source) in sources if source.subnetwork_id == subnetwork_id
     )
-    problem = allocation_problem(p, sources, capacity, subnetwork_id)
+    problem = allocation_problem(p_non_diff, sources, capacity, subnetwork_id)
     flow = JuMP.Containers.SparseAxisArray(Dict(only(problem[:F].axes) .=> 0.0))
 
     return AllocationModel(;