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

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 = S-infection
du = 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!

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:
# 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 = S-infection
du = 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)
``````

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.

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.