add possibility to call a direct implementation of the ODE RHS (#117)

* improve docstrings

* add possibility to call a direct implementation of the ODE RHS

* add tests

* format

* Apply suggestions from code review

Co-authored-by: Stefan Kopecz <[email protected]>

* more tests as suggested

---------

Co-authored-by: Stefan Kopecz <[email protected]>

Author: ranocha

Project.toml

@@ -1,7 +1,7 @@
 name = "PositiveIntegrators"
 uuid = "d1b20bf0-b083-4985-a874-dc5121669aa5"
 authors = ["Stefan Kopecz, Hendrik Ranocha, and contributors"]
-version = "0.2.2"
+version = "0.2.3"
 
 [deps]
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"

src/proddest.jl

@@ -3,13 +3,14 @@ abstract type AbstractPDSProblem end
 
 """
     PDSProblem(P, D, u0, tspan, p = NullParameters();
-                       p_prototype = nothing,
-                       analytic = nothing)
+               p_prototype = nothing,
+               analytic = nothing,
+               std_rhs = nothing)
 
 A structure describing a system of ordinary differential equations in form of a production-destruction system (PDS).
-`P` denotes the production matrix.
-The diagonal of `P` contains production terms without destruction counterparts.
-`D` is the vector of destruction terms without production counterparts.
+`P` denotes the function defining the production matrix ``P``.
+The diagonal of ``P`` contains production terms without destruction counterparts.
+`D` is the function defining the vector of destruction terms ``D`` without production counterparts.
 `u0` is the vector of initial conditions and `tspan` the time span
 `(t_initial, t_final)` of the problem. The optional argument `p` can be used
 to pass additional parameters to the functions `P` and `D`.
@@ -20,10 +21,16 @@ The functions `P` and `D` can be used either in the out-of-place form with signa
 ### Keyword arguments: ###
 
 - `p_prototype`: If `P` is given in in-place form, `p_prototype` or copies thereof are used to store evaluations of `P`.
-    If `p_prototype` is not specified explicitly and `P` is in-place, then `p_prototype` will be internally
+  If `p_prototype` is not specified explicitly and `P` is in-place, then `p_prototype` will be internally
   set to `zeros(eltype(u0), (length(u0), length(u0)))`.
 - `analytic`: The analytic solution of a PDS must be given in the form `f(u0,p,t)`.
-    Specifying the analytic solution can be useful for plotting and convergence tests.
+  Specifying the analytic solution can be useful for plotting and convergence tests.
+- `std_rhs`: The standard ODE right-hand side evaluation function callable
+  as `du = std_rhs(u, p, t)` for the out-of-place form and
+  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` and `D` is used.
 
 ## References
 
@@ -36,12 +43,13 @@ 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, Ta} <:
+struct PDSFunction{iip, specialize, P, D, PrototypeP, PrototypeD, StdRHS, Ta} <:
        AbstractODEFunction{iip}
     p::P
     d::D
     p_prototype::PrototypeP
     d_prototype::PrototypeD
+    std_rhs::StdRHS
     analytic::Ta
 end
 
@@ -82,18 +90,19 @@ end
 function PDSProblem{iip}(P, D, u0, tspan, p = NullParameters();
                          p_prototype = nothing,
                          analytic = nothing,
+                         std_rhs = 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)))
     end
-    # If a PDSFunction is to be evaluated and D is in-place, then d_prototype is used to store 
+    # If a PDSFunction is to be evaluated and D is in-place, then d_prototype is used to store
     # evaluations of D.
     d_prototype = similar(u0)
 
-    PD = PDSFunction{iip}(P, D; p_prototype = p_prototype, d_prototype = d_prototype,
-                          analytic = analytic)
+    PD = PDSFunction{iip}(P, D; p_prototype, d_prototype,
+                          analytic, std_rhs)
     PDSProblem{iip}(PD, u0, tspan, p; kwargs...)
 end
 
@@ -112,21 +121,45 @@ end
 # Most specific constructor for PDSFunction
 function PDSFunction{iip, FullSpecialize}(P, D;
                                           p_prototype = nothing,
-                                          d_prototype = nothing,
-                                          analytic = nothing) where {iip}
+                                          d_prototype,
+                                          analytic = nothing,
+                                          std_rhs = 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),
                 typeof(d_prototype),
-                typeof(analytic)}(P, D, p_prototype, d_prototype, analytic)
+                typeof(std_rhs), typeof(analytic)}(P, D, p_prototype, d_prototype, std_rhs,
+                                                   analytic)
 end
 
-# Evaluation of a PDSFunction (out-of-place)
+# Evaluation of a PDSFunction
 function (PD::PDSFunction)(u, p, t)
-    diag(PD.p(u, p, t)) + vec(sum(PD.p(u, p, t), dims = 2)) -
-    vec(sum(PD.p(u, p, t), dims = 1)) - vec(PD.d(u, p, t))
+    return PD.std_rhs(u, p, t)
 end
 
-# Evaluation of a PDSFunction (in-place)
 function (PD::PDSFunction)(du, u, p, t)
+    return PD.std_rhs(du, u, p, t)
+end
+
+# Default implementation of the standard right-hand side evaluation function
+struct PDSStdRHS{P, D, PrototypeP, PrototypeD} <: Function
+    p::P
+    d::D
+    p_prototype::PrototypeP
+    d_prototype::PrototypeD
+end
+
+# Evaluation of a PDSStdRHS (out-of-place)
+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)
+end
+
+# Evaluation of a PDSStdRHS (in-place)
+function (PD::PDSStdRHS)(du, u, p, t)
     PD.p(PD.p_prototype, u, p, t)
 
     if PD.p_prototype isa AbstractSparseMatrix
@@ -157,24 +190,32 @@ end
 """
     ConservativePDSProblem(P, u0, tspan, p = NullParameters();
                            p_prototype = nothing,
-                           analytic = nothing)
+                           analytic = nothing,
+                           std_rhs = nothing)
 
 A structure describing a conservative system of ordinary differential equation in form of a production-destruction system (PDS).
-`P` denotes the production matrix.
+`P` denotes the function defining the production matrix ``P``.
+The diagonal of ``P`` contains production terms without destruction counterparts.
 `u0` is the vector of initial conditions and `tspan` the time span
 `(t_initial, t_final)` of the problem. The optional argument `p` can be used
-to pass additional parameters to the function P.
+to pass additional parameters to the function `P`.
 
 The function `P` can be given either in the out-of-place form with signature
 `production_terms = P(u, p, t)` or the in-place form `P(production_terms, u, p, t)`.
 
 ### Keyword arguments: ###
 
 - `p_prototype`: If `P` is given in in-place form, `p_prototype` or copies thereof are used to store evaluations of `P`.
-    If `p_prototype` is not specified explicitly and `P` is in-place, then `p_prototype` will be internally
+  If `p_prototype` is not specified explicitly and `P` is in-place, then `p_prototype` will be internally
   set to `zeros(eltype(u0), (length(u0), length(u0)))`.
 - `analytic`: The analytic solution of a PDS must be given in the form `f(u0,p,t)`.
-    Specifying the analytic solution can be useful for plotting and convergence tests.
+  Specifying the analytic solution can be useful for plotting and convergence tests.
+- `std_rhs`: The standard ODE right-hand side evaluation function callable
+  as `du = std_rhs(u, p, t)` for the out-of-place form and
+  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.
 
 ## References
 
@@ -187,12 +228,12 @@ 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, TMP, Ta} <:
+struct ConservativePDSFunction{iip, specialize, P, PrototypeP, StdRHS, Ta} <:
        AbstractODEFunction{iip}
-    p::P
-    p_prototype::PrototypeP
-    tmp::TMP
-    analytic::Ta
+    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)
 end
 
 # define behavior of ConservativePDSFunction for non-existing fields
@@ -226,14 +267,15 @@ end
 function ConservativePDSProblem{iip}(P, u0, tspan, p = NullParameters();
                                      p_prototype = nothing,
                                      analytic = nothing,
+                                     std_rhs = 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)))
     end
 
-    PD = ConservativePDSFunction{iip}(P; p_prototype = p_prototype, analytic = analytic)
+    PD = ConservativePDSFunction{iip}(P; p_prototype, analytic, std_rhs)
     ConservativePDSProblem{iip}(PD, u0, tspan, p; kwargs...)
 end
 
@@ -252,18 +294,43 @@ end
 # Most specific constructor for ConservativePDSFunction
 function ConservativePDSFunction{iip, FullSpecialize}(P;
                                                       p_prototype = nothing,
-                                                      analytic = nothing) where {iip}
+                                                      analytic = nothing,
+                                                      std_rhs = 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)
+end
+
+# Evaluation of a ConservativePDSFunction
+function (PD::ConservativePDSFunction)(u, p, t)
+    return PD.std_rhs(u, p, t)
+end
+
+function (PD::ConservativePDSFunction)(du, u, p, t)
+    return PD.std_rhs(du, u, p, t)
+end
+
+# Default implementation of the standard right-hand side evaluation function
+struct ConservativePDSStdRHS{P, PrototypeP, TMP} <: Function
+    p::P
+    p_prototype::PrototypeP
+    tmp::TMP
+end
+
+function ConservativePDSStdRHS(P, p_prototype)
     if p_prototype isa AbstractSparseMatrix
         tmp = zeros(eltype(p_prototype), (size(p_prototype, 1),))
     else
         tmp = nothing
     end
-    ConservativePDSFunction{iip, FullSpecialize, typeof(P), typeof(p_prototype),
-                            typeof(tmp), typeof(analytic)}(P, p_prototype, tmp, analytic)
+    ConservativePDSStdRHS(P, p_prototype, tmp)
 end
 
-# Evaluation of a ConservativePDSFunction (out-of-place)
-function (PD::ConservativePDSFunction)(u, p, t)
+# Evaluation of a ConservativePDSStdRHS (out-of-place)
+function (PD::ConservativePDSStdRHS)(u, p, t)
     #vec(sum(PD.p(u, p, t), dims = 2)) - vec(sum(PD.p(u, p, t), dims = 1))
     P = PD.p(u, p, t)
 
@@ -277,7 +344,7 @@ function (PD::ConservativePDSFunction)(u, p, t)
     return f
 end
 
-function (PD::ConservativePDSFunction)(u::SVector, p, t)
+function (PD::ConservativePDSStdRHS)(u::SVector, p, t)
     P = PD.p(u, p, t)
 
     f = similar(u) #constructs MVector
@@ -296,8 +363,8 @@ function (PD::ConservativePDSFunction)(u::SVector, p, t)
     return SVector(f)
 end
 
-# Evaluation of a ConservativePDSFunction (in-place)
-function (PD::ConservativePDSFunction)(du, u, p, t)
+# Evaluation of a ConservativePDSStdRHS (in-place)
+function (PD::ConservativePDSStdRHS)(du, u, p, t)
     PD.p(PD.p_prototype, u, p, t)
     sum_terms!(du, PD.tmp, PD.p_prototype)
     return nothing

test/runtests.jl

@@ -216,6 +216,112 @@ end
         @test isnothing(check_no_stale_explicit_imports(PositiveIntegrators))
     end
 
+    @testset "ODE RHS" begin
+        let counter_p = Ref(1), counter_d = Ref(1), counter_rhs = Ref(1)
+            # out-of-place
+            prod1 = (u, p, t) -> begin
+                counter_p[] += 1
+                return [0 u[2]; u[1] 0]
+            end
+            dest1 = (u, p, t) -> begin
+                counter_d[] += 1
+                return zero(u)
+            end
+            rhs1 = (u, p, t) -> begin
+                counter_rhs[] += 1
+                return [-u[1] + u[2], u[1] - u[2]]
+            end
+            u0 = [1.0, 0.0]
+            tspan = (0.0, 1.0)
+            prob_default = PDSProblem(prod1, dest1, u0, tspan)
+            prob_special = PDSProblem(prod1, dest1, u0, tspan; std_rhs = rhs1)
+
+            counter_p[] = 0
+            counter_d[] = 0
+            counter_rhs[] = 0
+            @inferred prob_default.f(u0, nothing, 0.0)
+            @test counter_p[] == 1
+            @test counter_d[] == 1
+            @test counter_rhs[] == 0
+
+            counter_p[] = 0
+            counter_d[] = 0
+            counter_rhs[] = 0
+            @inferred prob_special.f(u0, nothing, 0.0)
+            @test counter_p[] == 0
+            @test counter_d[] == 0
+            @test counter_rhs[] == 1
+
+            # in-place
+            prod1! = (P, u, p, t) -> begin
+                counter_p[] += 1
+                P[1, 1] = 0
+                P[1, 2] = u[2]
+                P[2, 1] = u[1]
+                P[2, 2] = 0
+                return nothing
+            end
+            dest1! = (D, u, p, t) -> begin
+                counter_d[] += 1
+                fill!(D, 0)
+                return nothing
+            end
+            rhs1! = (du, u, p, t) -> begin
+                counter_rhs[] += 1
+                du[1] = -u[1] + u[2]
+                du[2] = u[1] - u[2]
+                return nothing
+            end
+            u0 = [1.0, 0.0]
+            tspan = (0.0, 1.0)
+            prob_default = PDSProblem(prod1!, dest1!, u0, tspan)
+            prob_special = PDSProblem(prod1!, dest1!, u0, tspan; std_rhs = rhs1!)
+
+            du = similar(u0)
+            counter_p[] = 0
+            counter_d[] = 0
+            counter_rhs[] = 0
+            @inferred prob_default.f(du, u0, nothing, 0.0)
+            @test counter_p[] == 1
+            @test counter_d[] == 1
+            @test counter_rhs[] == 0
+
+            counter_p[] = 0
+            counter_d[] = 0
+            counter_rhs[] = 0
+            @inferred prob_special.f(du, u0, nothing, 0.0)
+            @test counter_p[] == 0
+            @test counter_d[] == 0
+            @test counter_rhs[] == 1
+
+            counter_p[] = 0
+            counter_d[] = 0
+            counter_rhs[] = 0
+            @inferred solve(prob_default, MPE(); dt = 0.1)
+            @test 10 <= counter_p[] <= 11
+            @test 10 <= counter_d[] <= 11
+            @test counter_d[] == counter_p[]
+            @test counter_rhs[] == 0
+
+            counter_p[] = 0
+            counter_d[] = 0
+            counter_rhs[] = 0
+            @inferred solve(prob_default, Euler(); dt = 0.1)
+            @test 10 <= counter_p[] <= 11
+            @test 10 <= counter_d[] <= 11
+            @test counter_d[] == counter_p[]
+            @test counter_rhs[] == 0
+
+            counter_p[] = 0
+            counter_d[] = 0
+            counter_rhs[] = 0
+            @inferred solve(prob_special, Euler(); dt = 0.1)
+            @test counter_p[] == 0
+            @test counter_d[] == 0
+            @test 10 <= counter_rhs[] <= 11
+        end
+    end
+
     @testset "ConservativePDSFunction" begin
         prod_1! = (P, u, p, t) -> begin
             fill!(P, zero(eltype(P)))