Operators for Large-scale 3D simulations

In this documentation so far, we have talked about creating systems of ordinary differential equations in ModelingToolkit and then converting them to systems of partial differential equations to perform 1-, 2-, or 3-dimensional simulations. 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 Operator 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 processes that define behavior between grid cells.

ODE System

As an example, let's first define a system of ODEs:

using EarthSciMLBase
using ModelingToolkit, DifferentialEquations
using SciMLOperators, Plots
using DomainSets
using ModelingToolkit: t_nounits, D_nounits
t = t_nounits
D = D_nounits

@parameters y lon = 0.0 lat = 0.0 lev = 1.0 α = 10.0
@constants p = 1.0
@variables(
    u(t) = 1.0, v(t) = 1.0, x(t) = 1.0, y(t) = 1.0, windspeed(t) = 1.0
)

eqs = [D(u) ~ -α * √abs(v) + lon,
    D(v) ~ -α * √abs(u) + lat + 1e-14 * lev,
    windspeed ~ lat + lon + lev,
]
sys = ODESystem(eqs, t; name=:sys)

\[ \begin{align} \frac{\mathrm{d} u\left( t \right)}{\mathrm{d}t} &= \mathtt{lon} - \sqrt{\left|v\left( t \right)\right|} \alpha \\ \frac{\mathrm{d} v\left( t \right)}{\mathrm{d}t} &= \mathtt{lat} + 1 \cdot 10^{-14} \mathtt{lev} - \sqrt{\left|u\left( t \right)\right|} \alpha \\ \mathtt{windspeed}\left( t \right) &= \mathtt{lat} + \mathtt{lev} + \mathtt{lon} \end{align} \]

The equations above don't really have any physical meaning, but they include two state variables, some parameters, and a constant. There is also a variable windspeed which is "observed" based on the parameters, rather than being a state variable, which will be important later.

Operator

Next, we define an operator. To do so, first we create a new type that is a subtype of Operator:

struct ExampleOp <: Operator
end

Next, we need to define a method of EarthSciMLBase.get_scimlop for our operator. This method will be called to get a SciMLOperators.AbstractSciMLOperator that will be used conjunction with the ModelingToolkit system above to integrate the simulation forward in time.

function EarthSciMLBase.get_scimlop(op::ExampleOp, csys::CoupledSystem, mtk_sys, domain::DomainInfo, u0, p)
    α, trans1, trans2, trans3 = EarthSciMLBase.get_needed_vars(op, csys, mtk_sys, domain)

    obs_f! = ModelingToolkit.build_explicit_observed_function(mtk_sys,
        [α, trans1, trans2, trans3], checkbounds=false, return_inplace=true)[2]

    setp! = EarthSciMLBase.coord_setter(mtk_sys, domain)
    obscache = zeros(EarthSciMLBase.dtype(domain), 4)
    grd = EarthSciMLBase.grid(domain)

    function run(du, u, p, t)
        u = reshape(u, size(u0)...)
        du = reshape(du, size(u0)...)
        II = CartesianIndices(size(u)[2:end])
        for ix ∈ 1:size(u, 1)
            for I in II
                # Demonstrate coordinate transforms and observed values
                uu = view(u, :, I)
                setp!(p, I)
                obs_f!(obscache, u, p, t)
                t1, t2, t3, fv = obscache
                # Set derivative value.
                du[ix, I] = (t1 + t2 + t3) * fv
            end
        end
        nothing
    end
    FunctionOperator(run, u0[:], p=p)
end
nothing

The function above also doesn't have any physical meaning, but it demonstrates some functionality of the Operator "s". First, it retrieves a function to get the current value of an observed variable in our ODE system using the obs_functions argement, and it demonstrates how to call the resulting function to get that value. It also demonstrates how to get coordinate transforms using the coordinate_transform_functions argument. Coordinate transforms are discussed in more detail in the documentation for the DomainInfo type.

We also need to define a method of EarthSciMLBase.get_needed_vars, which will return which variables are needed by the operator.

function EarthSciMLBase.get_needed_vars(::ExampleOp, csys, mtk_sys, domain::DomainInfo)
    ts = EarthSciMLBase.partialderivative_transform_vars(mtk_sys, domain)
    return [mtk_sys.sys₊windspeed, ts...]
end
nothing

Domain

Once we have an ODE system and an operator, the final component we need is a domain to run the simulation on. Defining a domain is covered in more depth in the documentation for the DomainInfo type, but for now we'll just define a simple domain:

t_min = 0.0
lon_min, lon_max = -π, π
lat_min, lat_max = -0.45π, 0.45π
t_max = 11.5

indepdomain = t ∈ Interval(t_min, t_max)

partialdomains = [lon ∈ Interval(lon_min, lon_max),
    lat ∈ Interval(lat_min, lat_max),
    lev ∈ Interval(1, 3)]

domain = DomainInfo(
    partialderivatives_δxyδlonlat,
    constIC(16.0, indepdomain), constBC(16.0, partialdomains...);
    grid_spacing = [0.1π, 0.1π, 1])

Note that our domain includes a coordinate transform to convert from degrees latitude and longitude to meters. Our domain specification also includes grid spacing the the lon, lat, and lev coordinates, which we set as 0.1π, 0.1π, and 1, respectively.

Coupling and Running the Simulation

Next, initialize our operator, giving the the windspeed observed variable, and we can couple our ODESystem, Operator, and Domain together into a single model:

op = ExampleOp()

csys = couple(sys, op, domain)

CoupledSystem containing 1 system(s), 1 operator(s), and 0 callback(s).

Finally, we can choose a EarthSciMLBase.SolverStrategy and run the simulation. We choose the SolverStrangThreads strategy, which needs us to specify an ODE solver from the options available in DifferentialEquations.jl for both the MTK system. We choose the Tsit5 solver. Then we create an ODEProblem which can be used to run the simulation. Finally, we solve the problem using the solve function. At this point we need to choose a solver for the Operator part of the system, and we choose the Euler solver. We also choose a splitting time step of 1.0 seconds, which we pass both to our SolverStrangThreads strategy and to the solve function.

dt = 1.0 # Splitting time step
st = SolverStrangThreads(Tsit5(), 1.0)

prob = ODEProblem(csys, st)
sol = solve(prob, Euler(); dt=1.0)

After the simulation finishes, we can plot the result:

anim = @animate for i ∈ 1:length(sol.u)
    u = reshape(sol.u[i], :, size(domain)...)
    plot(
        heatmap(u[1, :, :, 1]),
        heatmap(u[1, :, :, 1]),
    )
end
gif(anim, fps = 15)