Merge pull request #50 from EarthSciML/strategy

Add SimulatorStrategies

Project.toml

@@ -1,22 +1,29 @@
 name = "EarthSciMLBase"
 uuid = "e53f1632-a13c-4728-9402-0c66d48804b0"
 authors = ["EarthSciML Authors and Contributors"]
-version = "0.12.2"
+version = "0.13.0"
 
 [deps]
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 Catalyst = "479239e8-5488-4da2-87a7-35f2df7eef83"
+DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
+BlockBandedMatrices = "0.13"
 Catalyst = "13"
+DiffEqCallbacks = "2"
 DifferentialEquations = "7"
 DocStringExtensions = "0.9"
 DomainSets = "0.5, 0.6"
@@ -26,6 +33,8 @@ MethodOfLines = "0.10"
 ModelingToolkit = "8"
 OrdinaryDiffEq = "6"
 ProgressLogging = "0.1"
+SciMLBase = "2"
+SciMLOperators = "0.3"
 Symbolics = "5"
 Unitful = "1"
 julia = "1.6"

docs/Project.toml

@@ -10,6 +10,7 @@ MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
 MethodOfLines = "94925ecb-adb7-4558-8ed8-f975c56a0bf4"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
@@ -24,5 +25,6 @@ MetaGraphsNext = "0.7"
 MethodOfLines = "0.8,0.9,0.10"
 ModelingToolkit = "8"
 Plots = "1"
+SciMLOperators = "0.3"
 Symbolics = "5"
 Unitful = "1"

docs/src/simulator.md

@@ -5,7 +5,7 @@ However, currently this does not work for large scale simulations.
 
 While this ModelingToolkit functionality is being built, we have a different solution based on the [`Simulator`](@ref) type in this package.
 Using this system, we still define systems of ODEs to define behavior in a single grid cell, and we also have [`Operator`](@ref) processes that define behavior between grid cells.
-The [`Simulator`](@ref) then integrates the ODEs and the Operators together using [Strang Splitting](https://en.wikipedia.org/wiki/Strang_splitting).
+The [`Simulator`](@ref) then integrates the ODEs and the Operators together.
 
 ## ODE System
 
@@ -14,7 +14,7 @@ As an example, let's first define a system of ODEs:
 ```@example sim
 using EarthSciMLBase
 using ModelingToolkit, DomainSets, DifferentialEquations
-using Plots
+using SciMLOperators, Plots
 
 @parameters y lon = 0.0 lat = 0.0 lev = 1.0 t α = 10.0
 @constants p = 1.0
@@ -35,7 +35,7 @@ There is also a variable `windspeed` which is "observed" based on the parameters
 
 ## Operator
 
-Next, we define an operator. To do so, first we create a new type that is a subtype of `Oerator`:
+Next, we define an operator. To do so, first we create a new type that is a subtype of [`Operator`](@ref):
 
 ```@example sim
 mutable struct ExampleOp <: Operator
@@ -44,27 +44,33 @@ end
 ```
 In the case above, we're setting up our operator so that it can hold a parameter from our ODE system.
 
-Next, we need to define a method of `EarthSciMLBase.run!` for our operator. This method will be called by the simulator to calculate the effect of the operator on the state variables.
+Next, we need to define a method of `EarthSciMLBase.get_scimlop` for our operator. This method will be called by the simulator to get a [`SciMLOperators.AbstractSciMLOperator`](https://docs.sciml.ai/SciMLOperators/stable/interface/) that will be used conjuction with the ModelingToolkit system above to integrate the simulation forward in time.
 
 ```@example sim
-function EarthSciMLBase.run!(op::ExampleOp, s::Simulator, t)
-    f = s.obs_fs[s.obs_fs_idx[op.α]]
-    for ix ∈ 1:size(s.u, 1)
-        for (i, c1) ∈ enumerate(s.grid[1])
-            for (j, c2) ∈ enumerate(s.grid[2])
-                for (k, c3) ∈ enumerate(s.grid[3])
-                    # Demonstrate coordinate transforms
-                    t1 = s.tf_fs[1](t, c1, c2, c3)
-                    t2 = s.tf_fs[2](t, c1, c2, c3)
-                    t3 = s.tf_fs[3](t, c1, c2, c3)
-                    # Demonstrate calculating observed value.
-                    fv = f(t, c1, c2, c3)
-                    # Set derivative value.
-                    s.du[ix, i, j, k] = (t1 + t2 + t3) * fv / 10
+function EarthSciMLBase.get_scimlop(op::ExampleOp, s::Simulator)
+    obs_f = s.obs_fs[s.obs_fs_idx[op.α]]
+    function run(du, u, p, t)
+        u = reshape(u, size(s.u)...)
+        du = reshape(du, size(s.u)...)
+        for ix ∈ 1:size(s.u, 1)
+            for (i, c1) ∈ enumerate(s.grid[1])
+                for (j, c2) ∈ enumerate(s.grid[2])
+                    for (k, c3) ∈ enumerate(s.grid[3])
+                        # Demonstrate coordinate transforms
+                        t1 = s.tf_fs[1](t, c1, c2, c3)
+                        t2 = s.tf_fs[2](t, c1, c2, c3)
+                        t3 = s.tf_fs[3](t, c1, c2, c3)
+                        # Demonstrate calculating observed value.
+                        fv = obs_f(t, c1, c2, c3)
+                        # Set derivative value.
+                        du[ix, i, j, k] = (t1 + t2 + t3) * fv
+                    end
                 end
             end
         end
+        nothing
     end
+    FunctionOperator(run, s.u[:], p=s.p)
 end
 ```
 The function above also doesn't have any physical meaning, but it demonstrates some functionality of the `Simulator` "`s`".
@@ -74,11 +80,6 @@ function to get that value.
 It also demonstrates how to get coordinate transforms using the `s.tf_fs` field.
 Coordinate transforms are discussed in more detail in the documentation for the [`DomainInfo`](@ref) type.
 
-Finally, we define a method of `EarthSciMLBase.timestep` for our operator. This method will be called by the simulator to determine the timestep for the operator.
-```@example sim
-EarthSciMLBase.timestep(op::ExampleOp) = 1.0
-```
-
 ## Domain
 
 Once we have an ODE system and an operator, the final component we need is a domain to run the simulation on.
@@ -122,22 +123,26 @@ nothing #hide
 
 ...and then create a Simulator. 
 Our simulator specification needs to include grid spacing the the `lon`, `lat`, and `lev`
-coordinates, which we set as 0.1π, 0.1π, and 1, respectively, and it needs us to 
-chose an ODE solver from the [options available in DifferentialEquations.jl](https://docs.sciml.ai/DiffEqDocs/stable/solvers/ode_solve/).
-We choose the `Tsit5` solver:
+coordinates, which we set as 0.1π, 0.1π, and 1, respectively.
 
 ```@example sim
-sim = Simulator(csys, [0.1π, 0.1π, 1], Tsit5())
+sim = Simulator(csys, [0.1π, 0.1π, 1])
 nothing #hide
 ```
 
-Finally, we can run the simulation:
+Finally, we can choose a [`EarthSciMLBase.SimulatorStrategy`](@ref) and run the simulation.
+We choose the [`SimulatorStrangThreads`](@ref) strategy, which needs us to 
+specify ODE solvers from the [options available in DifferentialEquations.jl](https://docs.sciml.ai/DiffEqDocs/stable/solvers/ode_solve/) for both the MTK system and the operator(s).
+We choose the `Tsit5` solver for our MTK system and the `Euler` solver for our operator.
+We also choose a time step of 1.0 seconds:
 
 ```@example sim
-run!(sim)
+st = SimulatorStrangThreads(Tsit5(), Euler(), 1.0)
+
+@time run!(sim, st)
 ```
 
-...and plot the result at the end of the simulation:
+After the simulation finishes, we can plot the result at the end of the simulation:
 
 ```@example sim
 plot(

src/EarthSciMLBase.jl

@@ -4,6 +4,9 @@ using Graphs, MetaGraphsNext
 using DocStringExtensions
 using Unitful
 using OrdinaryDiffEq, DomainSets
+using SciMLBase: DECallback, CallbackSet
+using DiffEqCallbacks
+using LinearAlgebra, BlockBandedMatrices
 using ProgressLogging
 
 include("add_dims.jl")
@@ -17,5 +20,7 @@ include("param_to_var.jl")
 include("graph.jl")
 include("simulator_utils.jl")
 include("simulator.jl")
+include("simulator_strategies.jl")
+include("simulator_strategy_strang.jl")
 
 end

src/coupled_system.jl

@@ -11,11 +11,31 @@ function systemhash(sys::ModelingToolkit.AbstractSystem)
     name
 end
 
+"""
+Types that implement an:
+
+`init_callback(x, Simulator)::DECallback`
+
+method can also be coupled into a `CoupledSystem`.
+The `init_callback` function will be run before the simulator is run
+to get the callback.
+"""
+init_callback() = error("Not implemented")
+
 """
 A system for composing together other systems using the [`couple`](@ref) function.
 
 $(FIELDS)
 
+Things that can be added to a `CoupledSystem`:
+    * `ModelingToolkit.ODESystem`s
+    * [`Operator`](@ref)s
+    * [`DomainInfo`](@ref)s
+    * [Callbacks](https://docs.sciml.ai/DiffEqDocs/stable/features/callback_functions/)
+    * Types `X` that implement a `EarthSciMLBase.get_callback(::X, ::Simulator)::DECallback` method
+    * Other `CoupledSystem`s
+    * Types `X` that implement a `EarthSciMLBase.couple(::X, ::CoupledSystem)` or `EarthSciMLBase.couple(::CoupledSystem, ::X)` method.
+    * `Tuple`s or `AbstractVector`s of any of the things above.
 """
 mutable struct CoupledSystem
     "Model components to be composed together"
@@ -32,10 +52,16 @@ mutable struct CoupledSystem
     A vector of operators to run during simulations.
     """
     ops::Vector{Operator}
+
+    "A vector of callbacks to run during simulations."
+    callbacks::Vector{DECallback}
+
+    "Objects `x` with an `init_callback(x, Simulator)::DECallback` method."
+    init_callbacks::Vector
 end
 
 function Base.show(io::IO, cs::CoupledSystem)
-    print(io, "CoupledSystem containing $(length(cs.systems)) system(s) and $(length(cs.ops)) operator(s).")
+    print(io, "CoupledSystem containing $(length(cs.systems)) system(s), $(length(cs.ops)) operator(s), and $(length(cs.callbacks) + length(cs.init_callbacks)) callback(s).")
 end
 
 """    
@@ -49,7 +75,7 @@ or any type `T` that has a method `couple(::CoupledSystem, ::T)::CoupledSystem`
 `couple(::T, ::CoupledSystem)::CoupledSystem` defined for it.
 """
 function couple(systems...)::CoupledSystem
-    o = CoupledSystem([], nothing, [], [])
+    o = CoupledSystem([], nothing, [], [], [], [])
     for sys ∈ systems
         if sys isa DomainInfo # Add domain information to the system.
             if o.domaininfo !== nothing
@@ -63,16 +89,25 @@ function couple(systems...)::CoupledSystem
         elseif sys isa CoupledSystem # Add a coupled system to the coupled system.
             o.systems = vcat(o.systems, sys.systems)
             o.pdefunctions = vcat(o.pdefunctions, sys.pdefunctions)
+            o.ops = vcat(o.ops, sys.ops)
+            o.callbacks = vcat(o.callbacks, sys.callbacks)
+            o.init_callbacks = vcat(o.init_callbacks, sys.init_callbacks)
             if sys.domaininfo !== nothing
                 if o.domaininfo !== nothing
                     error("Cannot add two sets of DomainInfo to a system.")
                 end
                 o.domaininfo = sys.domaininfo
             end
+        elseif sys isa DECallback
+            push!(o.callbacks, sys)
+        elseif hasmethod(init_callback, (typeof(sys), Simulator))
+            push!(o.init_callbacks, sys)
         elseif applicable(couple, o, sys)
             o = couple(o, sys)
         elseif applicable(couple, sys, o)
             o = couple(sys, o)
+        elseif (sys isa Tuple) || (sys isa AbstractVector)
+            o = couple(o, sys...)
         else
             error("Cannot couple a $(typeof(sys)).")
         end

src/operator.jl

@@ -4,24 +4,12 @@ export Operator
 Operators are objects that modify the current state of a `Simulator` system.
 Each operator should be define a `run` function with the signature:
 
-    `EarthSciMLBase.run!(op::Operator, s::Simulator, time, step_length)`
+    `EarthSciMLBase.get_scimlop(op::Operator, s::Simulator)::AbstractSciMLOperator`
 
-which should modify the `s.u` system state, and can also modify the `s.du` derivative cache if desired. 
-It should also implement:
-
-    `EarthSciMLBase.timestep(op::Operator)`
-
-which returns the timestep length for the operator.
-
-The Operator may also optionally implement the initialize! and finalize! methods,
-which will be run before the simulation starts and after it ends, respectively, if they are defined.
-
-    `EarthSciMLBase.initialize!(op::Operator, s::Simulator)`
-    `EarthSciMLBase.finalize!(op::Operator, s::Simulator)`
+which should return a [SciMLOperators.AbstractSciMLOperator](https://docs.sciml.ai/SciMLOperators/stable/interface/).
+Refer to the [SciMLOperators.jl](https://docs.sciml.ai/SciMLOperators/stable/)
+documentation for more information on how to define operators.
 """
 abstract type Operator end
 
-timestep(op::Operator) = ArgumentError("Operator $(typeof(op)) does not define a timestep function.")
-run!(op::Operator, _, _, _) = ArgumentError("Operator $(typeof(op)) does not define a run! function.")
-initialize!(_::Operator, _) = nothing
-finalize!(_::Operator, _) = nothing
+get_scimlop(op::Operator, _,) = ArgumentError("Operator $(typeof(op)) does not define a EarthSciMLBase.get_scimlop method.")