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

Specific Domains Modelling & Simulations
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)