[AI prototype] PositiveIntegrators.jl

.vscode/settings.json

@@ -69,6 +69,7 @@
         ".pixi"
     ],
     "julia.lint.run": true,
+    "julia.NumThreads": "1",
     "julia.numTestProcesses": 6,
     "notebook.codeActionsOnSave": {
         "notebook.source.fixAll": "explicit",

Project.toml

@@ -59,6 +59,7 @@ OteraEngine = "b2d7f28f-acd6-4007-8b26-bc27716e5513"
 PackageCompiler = "9b87118b-4619-50d2-8e1e-99f35a4d4d9d"
 PkgToSoftwareBOM = "6254a0f9-6143-4104-aa2e-fd339a2830a6"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+PositiveIntegrators = "d1b20bf0-b083-4985-a874-dc5121669aa5"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"

core/Project.toml

@@ -16,6 +16,7 @@ DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
@@ -33,6 +34,7 @@ MathOptAnalyzer = "d1179b25-476b-425c-b826-c7787f0fff83"
 MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
 Moshi = "2e0e35c7-a2e4-4343-998d-7ef72827ed2d"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
+NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 OrdinaryDiffEqBDF = "6ad6398a-0878-4a85-9266-38940aa047c8"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
 OrdinaryDiffEqDifferentiation = "4302a76b-040a-498a-8c04-15b101fed76b"
@@ -41,6 +43,7 @@ OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
 OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
 OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
+PositiveIntegrators = "d1b20bf0-b083-4985-a874-dc5121669aa5"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 SQLite = "0aa819cd-b072-5ff4-a722-6bc24af294d9"
@@ -64,11 +67,12 @@ DataInterpolations = "8.3"
 DataStructures = "0.18, 0.19"
 Dates = "1"
 DelimitedFiles = "1.9.1"
-DiffEqCallbacks = "3.6, 4"
-DifferentiationInterface = "0.6.48, 0.7"
+DiffEqBase = "7"
+DiffEqCallbacks = "4"
+DifferentiationInterface = "0.7"
 EnumX = "1.0"
 FiniteDiff = "2.21"
-ForwardDiff = "0.10.38, 1"
+ForwardDiff = "1"
 Graphs = "1.9"
 HiGHS = "1.7"
 IterTools = "1.4"
@@ -81,18 +85,19 @@ MathOptAnalyzer = "0.1.0"
 MetaGraphsNext = "0.6, 0.7, 0.8"
 Moshi = "0.3.7"
 NCDatasets = "0.14.8"
-OrdinaryDiffEqBDF = "1.14, 2"
-OrdinaryDiffEqCore = "3.2, 4"
-OrdinaryDiffEqDifferentiation = "2, 3"
-OrdinaryDiffEqLowOrderRK = "1.10, 2"
-OrdinaryDiffEqNonlinearSolve = "1.19, 2"
-OrdinaryDiffEqRosenbrock = "1.22, 2"
-OrdinaryDiffEqSDIRK = "1.11, 2"
-OrdinaryDiffEqTsit5 = "1.9, 2"
+NaNMath = "1"
+OrdinaryDiffEqBDF = "2.1.1"
+OrdinaryDiffEqCore = "4"
+OrdinaryDiffEqDifferentiation = "3"
+OrdinaryDiffEqLowOrderRK = "2"
+OrdinaryDiffEqNonlinearSolve = "2"
+OrdinaryDiffEqRosenbrock = "2"
+OrdinaryDiffEqSDIRK = "2"
+OrdinaryDiffEqTsit5 = "2"
 PrecompileTools = "1.2.1"
 Printf = "1"
 SQLite = "1.5.1"
-SciMLBase = "2.36, 3"
+SciMLBase = "3"
 SciMLOperators = "1.15.1"
 SparseArrays = "1"
 SparseConnectivityTracer = "1"

core/src/Ribasim.jl

@@ -32,13 +32,14 @@ using DifferentiationInterface:
 using ForwardDiff: derivative as forward_diff
 
 # Algorithms for solving ODEs.
-using OrdinaryDiffEqCore: OrdinaryDiffEqCore, get_du
-using OrdinaryDiffEqDifferentiation:
-    OrdinaryDiffEqDifferentiation, dolinsolve, jacobian2W!
-using SciMLOperators: WOperator, MatrixOperator
+using OrdinaryDiffEqCore: OrdinaryDiffEqCore, get_du, AbstractNLSolver
+using DiffEqBase: DiffEqBase, calculate_residuals!
+using OrdinaryDiffEqNonlinearSolve: OrdinaryDiffEqNonlinearSolve, relax!, _compute_rhs!
 import ADTypes
 using ADTypes: AutoForwardDiff
 import ForwardDiff
+import NaNMath
+import OrdinaryDiffEqBDF
 
 # Interface for defining and solving the ODE problem of the physical layer.
 using SciMLBase:
@@ -51,23 +52,21 @@ using SciMLBase:
     ODEProblem,
     get_proposed_dt,
     DEIntegrator,
-    FullSpecialize,
-    NoSpecialize,
-    SciMLOperators,
-    AbstractSciMLOperator,
-    LinearProblem,
-    LinearSolution
+    FullSpecialize
 
 # Automatically detecting the sparsity pattern of the Jacobian of water_balance!
 # through operator overloading
 using SparseConnectivityTracer: GradientTracer, TracerSparsityDetector
 using SparseMatrixColorings: GreedyColoringAlgorithm, sparsity_pattern
 
 # For efficient sparse computations
-using SparseArrays: SparseMatrixCSC, sparse, nzrange
+using SparseArrays: SparseMatrixCSC, sparse, nzrange, nonzeros
+
+# Positivity-preserving time integration
+using PositiveIntegrators: PDSProblem, MPRK22, MPRK43I, MPRK43II
 
 # Linear algebra
-using LinearAlgebra: LinearAlgebra, mul!, UniformScaling
+using LinearAlgebra: LinearAlgebra, mul!, UniformScaling, Symmetric, cholesky, Factorization
 
 # Interpolation functionality, used for e.g.
 # basin profiles and TabulatedRatingCurve. See also the node
@@ -182,7 +181,7 @@ include("allocation_init.jl")
 include("allocation_optim.jl")
 include("util.jl")
 include("graph.jl")
-include("differentiation.jl")
+include("pds.jl")
 include("model.jl")
 include("read.jl")
 include("write.jl")

core/src/allocation_optim.jl

@@ -14,7 +14,6 @@ function set_simulation_data!(
         user_demand,
         tabulated_rating_curve,
     ) = p.p_independent
-    du = get_du(integrator)
 
     errors = false
 
@@ -24,7 +23,7 @@ function set_simulation_data!(
     set_simulation_data!(allocation_model, linear_resistance, p, t)
     set_simulation_data!(allocation_model, manning_resistance, p, t)
     set_simulation_data!(allocation_model, tabulated_rating_curve, p, t)
-    set_simulation_data!(allocation_model, pump, outlet, du)
+    set_simulation_data!(allocation_model, pump, outlet, p.state_and_time_dependent_cache)
     set_simulation_data!(allocation_model, user_demand, t)
 
     if errors
@@ -350,7 +349,7 @@ function set_simulation_data!(
         allocation_model::AllocationModel,
         pump::Pump,
         outlet::Outlet,
-        du::CVector,
+        cache::StateAndTimeDependentCache,
     )::Nothing
     (; problem, scaling) = allocation_model
     pump_constraints = problem[:pump]
@@ -360,7 +359,7 @@ function set_simulation_data!(
     for node_id in only(pump_constraints.axes)
         constraint = pump_constraints[node_id]
         upstream_node_id = pump.inflow_link[node_id.idx].link[1]
-        q = du.pump[node_id.idx]
+        q = cache.current_flow_rate_pump[node_id.idx]
         if upstream_node_id.type == NodeType.Basin
             low_storage_factor = get_low_storage_factor(problem, upstream_node_id)
             JuMP.set_normalized_coefficient(
@@ -377,7 +376,7 @@ function set_simulation_data!(
     for node_id in only(outlet_constraints.axes)
         constraint = outlet_constraints[node_id]
         upstream_node_id = outlet.inflow_link[node_id.idx].link[1]
-        q = du.outlet[node_id.idx]
+        q = cache.current_flow_rate_outlet[node_id.idx]
         if upstream_node_id.type == NodeType.Basin
             low_storage_factor = get_low_storage_factor(problem, upstream_node_id)
             JuMP.set_normalized_coefficient(
@@ -769,19 +768,21 @@ function warm_start!(allocation_model::AllocationModel, integrator::DEIntegrator
     (; link_to_state_idx) = p.p_independent
 
     # Extrapolate the current instantaneous flow rates from the physical layer
+    # Flow rates are now stored in cache vectors, look up by link
+    flow_link_lookup = p.p_independent.graph[].flow_link_lookup
     for link in only(flow.axes)
-        state_index = get_state_index(getaxes(du), link_to_state_idx, link)
-        if !isnothing(state_index)
-            JuMP.set_start_value(flow[link], du[state_index] / scaling.flow)
+        link_idx = get_link_index(link, flow_link_lookup)
+        if !isnothing(link_idx)
+            JuMP.set_start_value(flow[link], p.p_independent.current_flow_rate[link_idx] / scaling.flow)
         end
     end
 
     # Extrapolate the current instantaneous storage rates from the physical layer
+    # du.basin[idx] now directly contains dS/dt
     for node_id in basin_ids_subnetwork
         JuMP.set_start_value(
             storage_change[node_id],
-            formulate_dstorage_wrt_time(du, p.p_independent, t, node_id) * Δt_allocation /
-                scaling.storage,
+            du.basin[node_id.idx] * Δt_allocation / scaling.storage,
         )
     end
 

core/src/bmi.jl

@@ -72,7 +72,7 @@ function BMI.get_value_ptr(model::Model, name::String)::Vector{Float64}
     elseif name == "basin.surface_runoff"
         basin.vertical_flux.surface_runoff::Vector{Float64}
     elseif name == "basin.cumulative_infiltration"
-        unsafe_array(u.infiltration)::Vector{Float64}
+        p_independent.cumulative_infiltration_total::Vector{Float64}
     elseif name == "basin.cumulative_drainage"
         basin.cumulative_drainage::Vector{Float64}
     elseif name == "basin.cumulative_surface_runoff"
@@ -82,7 +82,7 @@ function BMI.get_value_ptr(model::Model, name::String)::Vector{Float64}
     elseif name == "user_demand.demand"
         vec(user_demand.demand)::Vector{Float64}
     elseif name == "user_demand.cumulative_inflow"
-        unsafe_array(u.user_demand_inflow)::Vector{Float64}
+        state_and_time_dependent_cache.current_flow_rate_user_demand
     else
         error("Unknown variable $name")
     end