Tips on (sort of) creating a new Problem type in SciML, derived from an existing type

Specific Domains Modelling & Simulations
1

Hi everyone,

I’m trying to define what is in effect a new Problem in SciML, but what actually is a special case of an existing Problem. As I can’t derive from an instantiated type, I’m making a new type with an instance of the existing Problem as a member. Here’s my code, which is horrible, but works.The explanation is below, and I’m looking for suggestions on how to fit in more with SciML coding conventions.

import SciMLBase: @add_kwonly, DiscreteProblem
import OrdinaryDiffEq: solve, FunctionMap
import Random: AbstractRNG

struct MarkovProblem{uType, tType, P, F, K}
    f::F
    u0::uType
    tspan::Union{Tuple{tType,tType}, tType}
    p::P
    dt::tType
    rng::AbstractRNG
    prob::DiscreteProblem
    kwargs::K
    function MarkovProblem(f, u0, tspan, p, dt, rng; kwargs...)
        new{typeof(u0), typpeof(dt), typeof(p), typeof(f), typeof(kwargs)}(
                     f,
                     u0,
                     tspan,
                     p,
                     dt,
                     rng,
                     DiscreteProblem((du,u,p,t) -> f(du,u,p,t,dt,rng),
                                                   u0,
                                                   tspan,
                                                    p;
                                                    #dt=dt,
                                                    kwargs...)
                     )
    end

    @add_kwonly function MarkovProblem(f, u0, tspan, p, dt, rng; kwargs...)
        new{typeof(u0), typeof(dt), typeof(p), typeof(f), typeof(kwargs)}(
                     f,
                     u0,
                     tspan,
                     p,
                     dt,
                     rng,
                     DiscreteProblem((du,u,p,t) -> f(du,u,p,t,dt,rng),
                                                   u0,
                                                   tspan,
                                                    p;
                                                    #dt=dt,
                                                    kwargs...)
                     )
    end
end

function solve(prob::MarkovProblem)
    solve(prob.prob, FunctionMap(), dt=prob.dt)
end

function solve(prob::MarkovProblem, alg; kwargs...)
    solve(prob.prob, alg, dt=prob.dt)
end

## Test

using Random
using Distributions

function sir_markov!(du,u,p,t,dt,rng)
    (S,I) = u
    (β,γ) = p
    ifrac = 1-exp(-β*I*dt)
    rfrac = 1-exp(-γ*dt)
    infection=rand(rng, Binomial(S,ifrac))
    recovery=rand(rng, Binomial(I,rfrac))
    du[1] = S-infection
    du[2] = I+infection-recovery
    nothing
end

tspan = (0.0, 40.0)
dt = 0.1
u0 = [990, 10] # S,I
p = [0.0005, 0.25] # β,γ
rng = Xoshiro(1234)

## Solve direct using DiscreteProblem
discrete_prob = DiscreteProblem((du, u, p, t; kwargs...) -> sir_markov!(du,u,p,t,dt,rng; kwargs...),
                                 u0,tspan,p)
discrete_sol = solve(discrete_prob, FunctionMap(), dt=dt)

## Solve using MarkovProblem
markov_prob = MarkovProblem(sir_markov!, u0, tspan, p, dt, rng)
markov_sol = solve(markov_prob)

My aim is to be able to solve models where in addition to the usual ODEProblem signature of (du,u,p,t), I have (du,u,p,t,dt,rng), where dt is the time step, and rng is an AbstractRNG, so I can have per-simulation instances of random number generators with their own seed.

Here are some questions:

  1. What should solve (as well as init and step!) actually look like?
  2. How can I define a MarkovFunction class that fits with the current e.g. ODEFunction template?
  3. How can I best get a MarkovSolution that keeps all the standard solution-handling in SciML?
  4. Does my use of imports and using look OK?
  5. How can I specify one keyword argument (dt) but then be able to splat others (I commented out the line above)?
  6. How can my definition of MarkovProblem be improved?

The documentation in this area is a bit sparse, and I haven’t really delved deep into types in Julia before, so any pointers would be great!

Thanks!

2

Is this code cleaner? Is composition the best way here, rather than use traits e.g. via SimpleTraits.jl? I mostly need the types for the time and the state, so the typing should be sufficient. What are the scenarios for using the isinplace type?

# TODO:
# Add MarkovSystem
# Generalize to integer steps
# Write integrator struct
# Write init and step methods

import SciMLBase: DiscreteProblem, ODESolution
import OrdinaryDiffEq: solve, FunctionMap
import Random: AbstractRNG
import Plots: plot

struct MarkovProblem{uType, tType}
    prob::DiscreteProblem
    dt::tType
    rng::AbstractRNG
    function MarkovProblem(f, u0, tspan, p, dt, rng; kwargs...)
        prob = DiscreteProblem((du,u,p,t) -> f(du,u,p,t,dt,rng), u0, tspan, p; kwargs...)
        new{typeof(u0), typeof(dt)}(prob, dt, rng)
    end
end

struct MarkovSolution{uType, tType}
    sol::ODESolution
    function MarkovSolution(sol)
        new{eltype(eltype(sol.u)), eltype(sol.t)}(sol)
    end
end

function solve(prob::MarkovProblem; kwargs...)
    sol = solve(prob.prob, FunctionMap(), dt=prob.dt, kwargs...)
    MarkovSolution(sol)
end

function Base.show(io::IO, sol::MarkovSolution)
    A = sol.sol
    println(io, string("t: ", A.t))
    print(io, "u: ")
    show(io, A.u)
end

function Base.show(io::IO, m::MIME"text/plain", sol::MarkovSolution)
    A = sol.sol
    println(io, string("t: ", A.t))
    print(io, "u: ")
    show(io, m, A.u)
end

function plot(sol::MarkovSolution; kwargs...)
    plot(sol.sol; kwargs...)
end

## Test

using Random
using Distributions
using Plots

function sir_markov!(du,u,p,t,dt,rng)
    (S,I) = u
    (β,γ) = p
    ifrac = 1-exp(-β*I*dt)
    rfrac = 1-exp(-γ*dt)
    infection=rand(rng, Binomial(S,ifrac))
    recovery=rand(rng, Binomial(I,rfrac))
    du[1] = S-infection
    du[2] = I+infection-recovery
    nothing
end

tspan = (0.0, 40.0)
dt = 0.1
u0 = [990, 10] # S,I
p = [0.0005, 0.25] # β,γ
rng = Xoshiro(1234)

prob = MarkovProblem(sir_markov!, u0, tspan, p, dt, rng)
sol = solve(prob)
plot(sol)
3

Looks like a good start.

You need to implement a getindex function for your solution type, and a remake for your problem type but other than that this fulfils the basic interface. You may also want a getproperty for sol and prob that passes through to relevant fields in the wrapped type so more of the default functions would work.

I suggest that you put it in a module and make it subtype the relevant types in SciMLBase, with good enough docs it could be linked to from the SciMLDocs.

4

We generally don’t assume people will be making new problems outside of SciML so this interface is not defined. And we don’t plan to have it defined in a way that is so solidified for someone to do this any time soon. What you’ve done (other than the inference issues) looks fine, though there’s all of the symbolic indexing etc. that we really support as well, and that stuff is still evolving over time.