Store linear invariants matrix in (Conservative)PDSProblem

src/PDSProblemLibrary.jl

@@ -32,7 +32,8 @@ There is one independent linear invariant, e.g. ``u_1+u_2 = 1``.
   [DOI: 10.1016/S0168-9274(03)00101-6](https://doi.org/10.1016/S0168-9274(03)00101-6)
 """
 prob_pds_linmod = ConservativePDSProblem(P_linmod, u0_linmod, (0.0, 2.0),
-                                         analytic = f_linmod_analytic, std_rhs = f_linmod)
+                                         analytic = f_linmod_analytic, std_rhs = f_linmod,
+                                         linear_invariants = @SMatrix[1.0 1.0])
 
 function P_linmod!(P, u, p, t)
     P .= P_linmod(u, p, t)
@@ -51,7 +52,8 @@ Same as [`prob_pds_linmod`](@ref) but with in-place computation.
 prob_pds_linmod_inplace = ConservativePDSProblem(P_linmod!, Array(u0_linmod),
                                                  (0.0, 2.0),
                                                  analytic = f_linmod_analytic,
-                                                 std_rhs = f_linmod!)
+                                                 std_rhs = f_linmod!,
+                                                 linear_invariants = @SMatrix[1.0 1.0])
 
 # nonlinear model problem
 function P_nonlinmod(u, p, t)
@@ -86,7 +88,8 @@ There is one independent linear invariant, e.g. ``u_1+u_2+u_3 = 10.0``.
   [DOI: 10.1016/S0168-9274(03)00101-6](https://doi.org/10.1016/S0168-9274(03)00101-6)
 """
 prob_pds_nonlinmod = ConservativePDSProblem(P_nonlinmod, u0_nonlinmod, (0.0, 30.0),
-                                            std_rhs = f_nonlinmod)
+                                            std_rhs = f_nonlinmod,
+                                            linear_invariants = @SMatrix[1.0 1.0 1.0])
 
 # robertson problem
 function P_robertson(u, p, t)
@@ -119,7 +122,8 @@ There is one independent linear invariant, e.g. ``u_1+u_2+u_3 = 1.0``.
   2nd Edition, Springer (2002): Section IV.1.
 """
 prob_pds_robertson = ConservativePDSProblem(P_robertson, u0_robertson, (0.0, 1.0e11),
-                                            std_rhs = f_robertson)
+                                            std_rhs = f_robertson,
+                                            linear_invariants = @SMatrix[1.0 1.0 1.0])
 
 # brusselator problem
 function P_brusselator(u, p, t)
@@ -166,7 +170,9 @@ There are two independent linear invariants, e.g. ``u_1+u_4+u_5+u_6 = 10.2`` and
   [DOI: 10.1007/s10915-016-0267-9](https://doi.org/10.1007/s10915-016-0267-9)
 """
 prob_pds_brusselator = ConservativePDSProblem(P_brusselator, u0_brusselator, (0.0, 10.0),
-                                              std_rhs = f_brusselator)
+                                              std_rhs = f_brusselator,
+                                              linear_invariants = @SMatrix[1.0 0.0 0.0 1.0 1.0 1.0;
+                                                                           0.0 1.0 1.0 0.0 0.0 0.0])
 
 # SIR problem
 P_sir(u, p, t) = @SMatrix [0.0 0.0 0.0; 2*u[1]*u[2] 0.0 0.0; 0.0 u[2] 0.0]
@@ -195,7 +201,8 @@ There is one independent linear invariant, e.g. ``u_1+u_2+u_3 = 1.0``.
   Computers and Mathematics with Applications 66 (2013): 2307-2316.
   [DOI: 10.1016/j.camwa.2013.06.011](https://doi.org/10.1016/j.camwa.2013.06.011)
 """
-prob_pds_sir = ConservativePDSProblem(P_sir, u0_sir, (0.0, 20.0), std_rhs = f_sir)
+prob_pds_sir = ConservativePDSProblem(P_sir, u0_sir, (0.0, 20.0), std_rhs = f_sir,
+                                      linear_invariants = @SMatrix[1.0 1.0 1.0])
 
 # bertolazzi problem
 function P_bertolazzi(u, p, t)
@@ -237,7 +244,8 @@ There is one independent linear invariant, e.g. ``u_1+u_2+u_3 = 3.0``.
   [DOI: 10.1016/0898-1221(96)00142-3](https://doi.org/10.1016/0898-1221(96)00142-3)
 """
 prob_pds_bertolazzi = ConservativePDSProblem(P_bertolazzi, u0_bertolazzi, (0.0, 1.0),
-                                             std_rhs = f_bertolazzi)
+                                             std_rhs = f_bertolazzi,
+                                             linear_invariants = @SMatrix[1.0 1.0 1.0])
 
 # npzd problem
 function P_npzd(u, p, t)
@@ -288,7 +296,8 @@ There is one independent linear invariant, e.g. ``u_1+u_2+u_3+u_4 = 15.0``.
   Ocean Dynamics 55 (2005): 326-337.
   [DOI: 10.1007/s10236-005-0001-x](https://doi.org/10.1007/s10236-005-0001-x)
 """
-prob_pds_npzd = ConservativePDSProblem(P_npzd, u0_npzd, (0.0, 10.0), std_rhs = f_npzd)
+prob_pds_npzd = ConservativePDSProblem(P_npzd, u0_npzd, (0.0, 10.0), std_rhs = f_npzd,
+                                       linear_invariants = @SMatrix[1.0 1.0 1.0 1.0])
 
 # stratospheric reaction problem
 function P_stratreac(u, p, t)
@@ -437,7 +446,9 @@ There are two independent linear invariants, e.g. ``u_1+u_2+3u_3+2u_4+u_5+2u_6=(
   [DOI: 10.2140/camcos.2021.16.155](https://doi.org/10.2140/camcos.2021.16.155)
 """
 prob_pds_stratreac = PDSProblem(P_stratreac, d_stratreac, u0_stratreac, (4.32e4, 3.024e5),
-                                std_rhs = f_stratreac)
+                                std_rhs = f_stratreac,
+                                linear_invariants = @SMatrix[1.0 1.0 3.0 2.0 1.0 2.0;
+                                                             0.0 0.0 0.0 0.0 1.0 1.0])
 
 # mapk problem
 function f_minmapk(u, p, t)
@@ -518,4 +529,6 @@ There are two independent linear invariants, e.g. ``u_1+u_4+u_6=1.75`` and ``u_2
   PLoS ONE 12 (2017): e0178457.
   [DOI: 10.1371/journal.pone.0178457](https://doi.org/10.1371/journal.pone.0178457)
 """
-prob_pds_minmapk = PDSProblem(P_minmapk, D_minmapk, u0, tspan; std_rhs = f_minmapk)
+prob_pds_minmapk = PDSProblem(P_minmapk, D_minmapk, u0, tspan; std_rhs = f_minmapk,
+                              linear_invariants = @SMatrix[1.0 0.0 0.0 1.0 0.0 1.0;
+                                                           0.0 1.0 1.0 1.0 1.0 0.0])

src/proddest.jl

@@ -5,7 +5,8 @@ abstract type AbstractPDSProblem end
     PDSProblem(P, D, u0, tspan, p = NullParameters();
                p_prototype = nothing,
                analytic = nothing,
-               std_rhs = nothing)
+               std_rhs = nothing,
+               linear_invariants = nothing)
 
 A structure describing a system of ordinary differential equations in form of a production-destruction system (PDS).
 `P` denotes the function defining the production matrix ``P``.
@@ -31,6 +32,9 @@ The functions `P` and `D` can be used either in the out-of-place form with signa
   the production-destruction representation of the ODE, will use this function
   instead to compute the solution. If not specified,
   a default implementation calling `P` and `D` is used.
+-`linear_invariants`: The rows of this matrix contain the linear invariants of the ODE. 
+  Certain solvers or callbacks require this matrix.
+  Note that this feature is experimental and its API may change in future releases.
 
 ## References
 
@@ -43,14 +47,15 @@ The functions `P` and `D` can be used either in the out-of-place form with signa
 struct PDSProblem{iip} <: AbstractPDSProblem end
 
 # New ODE function PDSFunction
-struct PDSFunction{iip, specialize, P, D, PrototypeP, PrototypeD, StdRHS, Ta} <:
+struct PDSFunction{iip, specialize, P, D, PrototypeP, PrototypeD, StdRHS, Ta, LI} <:
        AbstractODEFunction{iip}
     p::P
     d::D
     p_prototype::PrototypeP
     d_prototype::PrototypeD
     std_rhs::StdRHS
     analytic::Ta
+    linear_invariants::LI
 end
 
 # define behavior of PDSFunctions for non-existing fields
@@ -73,8 +78,7 @@ function Base.getproperty(obj::PDSFunction, sym::Symbol)
 end
 
 # Most general constructor for PDSProblems
-function PDSProblem(P, D, u0, tspan, p = NullParameters();
-                    kwargs...)
+function PDSProblem(P, D, u0, tspan, p = NullParameters(); kwargs...)
     Piip = isinplace(P, 4)
     Diip = isinplace(D, 4)
     if Piip == Diip
@@ -87,11 +91,16 @@ end
 
 # Specialized constructor for PDSProblems setting `iip` manually
 # (arbitrary functions)
-function PDSProblem{iip}(P, D, u0, tspan, p = NullParameters();
+function PDSProblem{iip}(P,
+                         D,
+                         u0,
+                         tspan,
+                         p = NullParameters();
                          p_prototype = nothing,
                          analytic = nothing,
                          std_rhs = nothing,
-                         kwargs...) where {iip}
+                         linear_invariants = nothing,
+                         kwargs...,) where {iip}
 
     # p_prototype is used to store evaluations of P, if P is in-place.
     if isnothing(p_prototype) && iip
@@ -101,8 +110,8 @@ function PDSProblem{iip}(P, D, u0, tspan, p = NullParameters();
     # evaluations of D.
     d_prototype = similar(u0 ./ oneunit(first(tspan)))
 
-    PD = PDSFunction{iip}(P, D; p_prototype, d_prototype,
-                          analytic, std_rhs)
+    PD = PDSFunction{iip}(P, D; p_prototype, d_prototype, analytic, std_rhs,
+                          linear_invariants)
     PDSProblem{iip}(PD, u0, tspan, p; kwargs...)
 end
 
@@ -119,18 +128,26 @@ function PDSFunction{iip}(P, D; kwargs...) where {iip}
 end
 
 # Most specific constructor for PDSFunction
-function PDSFunction{iip, FullSpecialize}(P, D;
+function PDSFunction{iip, FullSpecialize}(P,
+                                          D;
                                           p_prototype = nothing,
                                           d_prototype = nothing,
                                           analytic = nothing,
-                                          std_rhs = nothing) where {iip}
+                                          std_rhs = nothing,
+                                          linear_invariants = nothing,) where {iip}
     if std_rhs === nothing
         std_rhs = PDSStdRHS(P, D, p_prototype, d_prototype)
     end
-    PDSFunction{iip, FullSpecialize, typeof(P), typeof(D), typeof(p_prototype),
+    PDSFunction{iip,
+                FullSpecialize,
+                typeof(P),
+                typeof(D),
+                typeof(p_prototype),
                 typeof(d_prototype),
-                typeof(std_rhs), typeof(analytic)}(P, D, p_prototype, d_prototype, std_rhs,
-                                                   analytic)
+                typeof(std_rhs),
+                typeof(analytic),
+                typeof(linear_invariants)}(P, D, p_prototype, d_prototype, std_rhs,
+                                           analytic, linear_invariants)
 end
 
 # Evaluation of a PDSFunction
@@ -165,8 +182,7 @@ end
 function (PD::PDSStdRHS)(u, p, t)
     P = PD.p(u, p, t)
     D = PD.d(u, p, t)
-    diag(P) + vec(sum(P, dims = 2)) -
-    vec(sum(P, dims = 1)) - vec(D)
+    diag(P) + vec(sum(P; dims = 2)) - vec(sum(P; dims = 1)) - vec(D)
 end
 
 # Evaluation of a PDSStdRHS (in-place)
@@ -203,7 +219,8 @@ end
     ConservativePDSProblem(P, u0, tspan, p = NullParameters();
                            p_prototype = nothing,
                            analytic = nothing,
-                           std_rhs = nothing)
+                           std_rhs = nothing,
+                           linear_invariants = nothing)
 
 A structure describing a conservative system of ordinary differential equation in form of a production-destruction system (PDS).
 `P` denotes the function defining the production matrix ``P``.
@@ -227,7 +244,10 @@ The function `P` can be given either in the out-of-place form with signature
   as `std_rhs(du, u, p, t)` for the in-place form. Solvers that do not rely on
   the production-destruction representation of the ODE, will use this function
   instead to compute the solution. If not specified,
-  a default implementation calling `P` is used.
+  a default implementation calling `P` is used
+-`linear_invariants`: The rows of this matrix contain the linear invariants of the ODE. 
+  Certain solvers or callbacks require this matrix.
+  Note that this feature is experimental and its API may change in future releases.
 
 ## References
 
@@ -240,12 +260,13 @@ The function `P` can be given either in the out-of-place form with signature
 struct ConservativePDSProblem{iip} <: AbstractPDSProblem end
 
 # New ODE function ConservativePDSFunction
-struct ConservativePDSFunction{iip, specialize, P, PrototypeP, StdRHS, Ta} <:
+struct ConservativePDSFunction{iip, specialize, P, PrototypeP, StdRHS, Ta, LI} <:
        AbstractODEFunction{iip}
     p::P # production terms
     p_prototype::PrototypeP # prototype for production terms
     std_rhs::StdRHS # standard right-hand side evaluation function
     analytic::Ta # analytic solution (or nothing)
+    linear_invariants::LI
 end
 
 # define behavior of ConservativePDSFunction for non-existing fields
@@ -268,26 +289,29 @@ function Base.getproperty(obj::ConservativePDSFunction, sym::Symbol)
 end
 
 # Most general constructor for ConservativePDSProblems
-function ConservativePDSProblem(P, u0, tspan, p = NullParameters();
-                                kwargs...)
+function ConservativePDSProblem(P, u0, tspan, p = NullParameters(); kwargs...)
     iip = isinplace(P, 4)
     return ConservativePDSProblem{iip}(P, u0, tspan, p; kwargs...)
 end
 
 # Specialized constructor for ConservativePDSProblems setting `iip` manually
-# (arbitrary function)
-function ConservativePDSProblem{iip}(P, u0, tspan, p = NullParameters();
+# (arbitrary function)  
+function ConservativePDSProblem{iip}(P,
+                                     u0,
+                                     tspan,
+                                     p = NullParameters();
                                      p_prototype = nothing,
                                      analytic = nothing,
                                      std_rhs = nothing,
-                                     kwargs...) where {iip}
+                                     linear_invariants = nothing,
+                                     kwargs...,) where {iip}
 
     # p_prototype is used to store evaluations of P, if P is in-place.
     if isnothing(p_prototype) && iip
         p_prototype = zeros(eltype(u0), (length(u0), length(u0))) / oneunit(first(tspan))
     end
 
-    PD = ConservativePDSFunction{iip}(P; p_prototype, analytic, std_rhs)
+    PD = ConservativePDSFunction{iip}(P; p_prototype, analytic, std_rhs, linear_invariants)
     ConservativePDSProblem{iip}(PD, u0, tspan, p; kwargs...)
 end
 
@@ -304,16 +328,20 @@ function ConservativePDSFunction{iip}(P; kwargs...) where {iip}
 end
 
 # Most specific constructor for ConservativePDSFunction
-function ConservativePDSFunction{iip, FullSpecialize}(P;
-                                                      p_prototype = nothing,
-                                                      analytic = nothing,
-                                                      std_rhs = nothing) where {iip}
+function ConservativePDSFunction{iip, FullSpecialize}(P; p_prototype = nothing,
+                                                      analytic = nothing, std_rhs = nothing,
+                                                      linear_invariants = nothing) where {iip}
     if std_rhs === nothing
         std_rhs = ConservativePDSStdRHS(P, p_prototype)
     end
-    ConservativePDSFunction{iip, FullSpecialize, typeof(P), typeof(p_prototype),
-                            typeof(std_rhs), typeof(analytic)}(P, p_prototype, std_rhs,
-                                                               analytic)
+    ConservativePDSFunction{iip,
+                            FullSpecialize,
+                            typeof(P),
+                            typeof(p_prototype),
+                            typeof(std_rhs),
+                            typeof(analytic),
+                            typeof(linear_invariants)}(P, p_prototype, std_rhs, analytic,
+                                                       linear_invariants)
 end
 
 # Evaluation of a ConservativePDSFunction

test/runtests.jl

@@ -4,7 +4,7 @@ using SparseArrays
 using Statistics: mean, median
 
 using DoubleFloats: Double64
-using StaticArrays: SMatrix, MVector, @SVector, SVector, SA
+using StaticArrays: @SMatrix, SMatrix, MVector, @SVector, SVector, SA
 
 using Unitful: @u_str, ustrip
 
@@ -2680,4 +2680,70 @@ end
             end
         end
     end
+    @testset "linear invariant matrices" begin
+        @testset "Check that the structure of the linear invariant field is well implemented" begin
+            u0 = [0.9, 0.1]
+            tspan = (0.0, 2.0)
+            AT = @SMatrix[1 1]
+
+            linmodP(u, p, t) = [0 u[2]; 5*u[1] 0]
+            linmodD(u, p, t) = [0.0; 0.0]
+            linmod_PDS_op = PDSProblem(linmodP, linmodD, u0, tspan, linear_invariants = AT)
+
+            linmod_PDS_op_cons = ConservativePDSProblem(linmodP, linmodD, u0, tspan,
+                                                        linear_invariants = AT)
+
+            @test linmod_PDS_op.f.linear_invariants == AT
+            @test linmod_PDS_op_cons.f.linear_invariants == AT
+        end
+
+        @testset "Check that the predefined linear invariant matrices fit in predefined problems" begin
+            # non-stiff conservative problems (out-of-place)
+            probs = (prob_pds_linmod, prob_pds_nonlinmod, prob_pds_brusselator,
+                     prob_pds_sir, prob_pds_npzd)
+            @testset "$prob" for prob in probs
+                dt = (last(prob.tspan) - first(prob.tspan)) / 1e4
+                sol = solve(prob, Tsit5(); dt, isoutofdomain = isnegative)
+                @test prob.f.linear_invariants * prob.u0 ≈
+                      prob.f.linear_invariants * sol.u[end]
+            end
+
+            # non-stiff conservative problems (in-place)
+            probs = (prob_pds_linmod_inplace,)
+            @testset "$prob" for prob in probs
+                dt = (last(prob.tspan) - first(prob.tspan)) / 1e4
+                sol = solve(prob, Tsit5(); dt, isoutofdomain = isnegative)
+                @test prob.f.linear_invariants * prob.u0 ≈
+                      prob.f.linear_invariants * sol.u[end]
+            end
+
+            # non-stiff non-conservative problems (out-of-place)
+            probs = (prob_pds_minmapk,)
+            @testset "$prob" for prob in probs
+                dt = (last(prob.tspan) - first(prob.tspan)) / 1e4
+                sol = solve(prob, Tsit5(); dt, isoutofdomain = isnegative)
+                @test prob.f.linear_invariants * prob.u0 ≈
+                      prob.f.linear_invariants * sol.u[end]
+            end
+
+            # Robertson problem
+            prob = prob_pds_robertson
+            dt = 1e-6
+            sol = solve(prob, Rosenbrock23(); dt)
+            @test prob.f.linear_invariants * prob.u0 ≈ prob.f.linear_invariants * sol.u[end]
+
+            # Bertolazzi problem
+            prob = prob_pds_bertolazzi
+            alg = ImplicitEuler()
+            dt = 1e-3
+            sol = solve(prob, alg; dt)
+            @test prob.f.linear_invariants * prob.u0 ≈ prob.f.linear_invariants * sol.u[end]
+
+            # Stratospheric reaction problem
+            prob = prob_pds_stratreac
+            dt = 1.0
+            sol = solve(prob, Rosenbrock23(); dt)
+            @test prob.f.linear_invariants * prob.u0 ≈ prob.f.linear_invariants * sol.u[end]
+        end
+    end
 end;