mschauer / mitosisstochasticdiffeq.jl Goto Github PK

Backward-filtering forward-guiding with StochasticDiffEq.jl

License: MIT License

Julia 100.00%

mitosisstochasticdiffeq.jl's Introduction

MitosisStochasticDiffEq.jl

Implements the Mitosis transformation rules backwardfilter and forwardguiding for SciML's StochasticDiffEq problems.

If the (possibly non-linear) drift depends linearly on parameters, estimate the parameters from continuous observations by regression.

Synopsis

MitosisStochasticDiffEq implements the backward filter and the forward change of measure of the Automatic Backward Filtering Forward Guiding paradigm (van der Meulen and Schauer, 2020) as transformation rules for SDE models, suitable to be incorporated into probabilistic programming approaches.

In particular, this package implements the equations ... of section 9.1, [2] further detailed in [1]. The recursion for the quantity c in [1, Theorem 3.3 (Information filter)] is replaced by the simpler rule from [2, Example 10.8.]

Show reel

Bayesian regression on the drift parameter of an SDE

using StochasticDiffEq
using Random
using MitosisStochasticDiffEq
import MitosisStochasticDiffEq as MSDE
using LinearAlgebra, Statistics

# Model and sensitivity
function f(du, u, θ, t)
    c = 0.2 * θ
    du[1] = -0.1 * u[1] + c * u[2]
    du[2] = - c * u[1] - 0.1 * u[2]
    return
end
function g(du, u, θ, t)
    fill!(du, 0.15)
    return
end

# b is linear in the parameter with Jacobian
function b_jac(J,x,θ,t)
    J .= false
    J[1,1] =   0.2 * x[2]
    J[2,1] = - 0.2 * x[1]
    nothing
end
# and intercept
function b_icpt(dx,x,θ,t)
    dx .= false
    dx[1] = -0.1 * x[1]
    dx[2] = -0.1 * x[2]
    nothing
end

# Simulate path ensemble
x0 = [1.0, 1.0]
tspan = (0.0, 20.0)
θ0 = 1.0
dt = 0.05
t = range(tspan...; step=dt)

prob = SDEProblem{true}(f, g, x0, tspan, θ0)
ensembleprob = EnsembleProblem(prob)
ensemblesol = solve(
    ensembleprob, EM(), EnsembleThreads(); dt=dt, saveat=t, trajectories=1000
)

# Inference on drift parameters
sdekernel = MSDE.SDEKernel(f,g,t,0*θ0)
ϕprototype = zeros((length(x0),length(θ0))) # prototypes for vectors
yprototype = zeros((length(x0),))
R = MSDE.Regression!(sdekernel,yprototype,paramjac_prototype=ϕprototype,paramjac=b_jac,intercept=b_icpt)
prior_precision = 0.1I(1)
posterior = MSDE.conjugate(R, ensemblesol, prior_precision)
print(mean(posterior)[], " ± ", sqrt(cov(posterior)[]))

References

[1] Marcin Mider, Moritz Schauer, Frank van der Meulen (2020): Continuous-discrete smoothing of diffusions. [arxiv:1712.03807].
[2] Frank van der Meulen, Moritz Schauer (2020): Automatic Backward Filtering Forward Guiding for Markov processes and graphical models. [arXiv:2010.03509].

mitosisstochasticdiffeq.jl's People

Contributors

Stargazers

Watchers

Forkers

pitmonticone

mitosisstochasticdiffeq.jl's Issues

size SArray not known

I tried running bridgeexample.jl

solfw, ll = MitosisStochasticDiffEq.forwardguiding(sdekernel, message, (u0, 0.0),
   Z=nothing; inplace=false, save_noise=true)

This gives an error

julia> solfw, ll = MitosisStochasticDiffEq.forwardguiding(sdekernel, message, (u0, 0.0),
          Z=nothing; inplace=false, save_noise=true)
ERROR: The size of type `SArray` is not known.

If you were trying to construct (or `convert` to) a `StaticArray` you
may need to add the size explicitly as a type parameter so its size is
inferrable to the Julia compiler (or performance would be terrible). For
example, you might try

    m = zeros(3,3)
    SMatrix(m)            # this error
    SMatrix{3,3}(m)       # correct - size is inferrable
    SArray{Tuple{3,3}}(m) # correct, note Tuple{3,3}

Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] missing_size_error(::Type{SArray}) at /Users/frankvandermeulen/.julia/packages/StaticArrays/LJQEe/src/traits.jl:72
 [3] Size(::Type{SArray}) at /Users/frankvandermeulen/.julia/packages/StaticArrays/LJQEe/src/traits.jl:88
 [4] length(::Type{SArray}) at /Users/frankvandermeulen/.julia/packages/StaticArrays/LJQEe/src/abstractarray.jl:2
 [5] convert at /Users/frankvandermeulen/.julia/packages/StaticArrays/LJQEe/src/convert.jl:28 [inlined]
 [6] StaticArray at /Users/frankvandermeulen/.julia/packages/StaticArrays/LJQEe/src/convert.jl:7 [inlined]
 [7] mypack(::SArray{Tuple{4},Float64,1,4}, ::Float64) at /Users/frankvandermeulen/.julia/packages/MitosisStochasticDiffEq/iVjkQ/src/guiding.jl:16
 [8] forwardguiding(::MitosisStochasticDiffEq.SDEKernel{typeof(f),typeof(g),StepRangeLen{Float64,Base.TwicePrecision{Float64},B

GuidingDiffusionCache

Should the final line be as below? (the loglikelihood is always scalar valued)

function (G::GuidingDiffusionCache)(u,p,t)
  @unpack g = G

  x = @view u[1:end-1]
  dx = g(x,p,t)
  return mypack(dx, 0.0) #[dx; zero(eltype(u))]
end

flag constant diffusivity

Computing the loglikelihood is much cheaper if it is know that the diffusivity is constant. In that case, the second term on the RHS

dl = dot(f(x,p,t) -  ktilde.f(x,ktilde.p,t), r) - 0.5*tr((outer_(g(x,p,t)) - outer_(ktilde.g(x,ktilde.p,t)))*(inv(Σ) - outer_(r)))

need not be computed. One option could be to to have a field in the SDEkernel specifying this. A check that indeed the diffusivities of k and ktilde are the same would be nice.

[Also node that evaluation of the first term on the RHS is not type-stable with staticvectors.. this issue has been reported.]

Allow for time-dependence of tilde kernel parameters in `AffineMap` and `ConstantMap`

In order to not lose our Slack conversation @mschauer:

#34 restructures the SDE kernel.

Instead of AffineMap, we could define and overload TAffineMap
with TAffineMap(t, x, p) and B, β, σ̃ = linearparams(f, g, t) such that the tilde kernel can be defined as

kernel = SDEKernel(TAffineMap, ConstantMap, trange, (B, β, σ̃))

to support time-dependence.

Fix the test taking the unregistered dependence Mitosis into account.

Perhaps @mmider knows how to fix this?

Todo list

Now that the test show that this is all fine 🎉, some next steps

Guiding function for H,F parametrization

#48 implemented the information filter (H,F parametrization). We still need to add the associated guiding functions in this parametrization and finish the convert functions to convert messages between the H,F and the \nu,P parametrization.

Stacking messages together

We compute the messages between observations. For a single observation at the endpoint, we thus have a single message. In the case of several observations over time, we would like to stack the individual messages together such that we are back at a single longer message (nu,P) over the full time span of the forward model.

Lyapunov filter

#48 contains some first structure/code snippets for the Lyapunov filter (Proposition 3.5. in https://arxiv.org/abs/1712.03807) when B and \tilde{a} are constant.

Based on the discussions in #48, there are a few points remaining:

Potentially there is still a small error in the solution for \nu because the solution from the Lyapunov filter doesn't agree with the covariance filter based on solving the ODE yet.
a time-inhomogeneous example should be added. There is https://github.com/mschauer/Bridge.jl/blob/430e9d8a0a39a2cd00f13cd10bfd1f3df5f3d2c4/src/linpro.jl#L13 in Bridge.jl, where a time dependent drift using a cspline (see https://github.com/mschauer/Bridge.jl/blob/430e9d8a0a39a2cd00f13cd10bfd1f3df5f3d2c4/src/cspline.jl
and https://en.wikipedia.org/wiki/Cubic_Hermite_spline) is defined.

(lack of) consistency in parametrisation

I see both

(c, mu, Sigma)
(c, nu, P)

I guess these are meant to be the same. With all notations, shall we try to stick exactly to the notation in the paper Mider/Schauer/vdMeulen?

Register 0.x

Let’s register sometimes

mypack with SArray

Rather than mypack(a::SArray,c::Number) = SArray([a; c]) should be

mypack(a::SArray,c::Number) = SVector{length(a)+1}([a; c])

(don't know if this is efficient...)

TagBot trigger issue

This issue is used to trigger TagBot; feel free to unsubscribe.

If you haven't already, you should update your TagBot.yml to include issue comment triggers.
Please see this post on Discourse for instructions and more details.

If you'd like for me to do this for you, comment TagBot fix on this issue.
I'll open a PR within a few hours, please be patient!

Replace tstart, tend and dt by range

and allow non-equidistant time grids.

If the observation noise is small, it is a good idea to transform the time grid by

timechange(s, tstart=first(s), tend=last(s)) = tstart .+ (s .- tstart).*(2.0 .- (s .- tstart)/(tend-tstart))

Allow StaticArrays for θlin

Currently, we get

ERROR: MethodError: Cannot `convert` an object of type 
  ArrayPartition{Float64{},Tuple{SVector{2, Float64},Array{Float64{},2},Array{Float64{},1}}} to an object of type 
  ArrayPartition{Float64{},Tuple{Vector{Float64},Array{Float64{},2},Array{Float64{},1}}}
Closest candidates are:
  convert(::Type{T}, ::Intervals.Interval{T, L, R} where {L<:Intervals.Bound, R<:Intervals.Bound}) where T at /Users/smoritz/.julia/packages/Intervals/ua9cq/src/interval.jl:253
  convert(::Type{T}, ::Intervals.AnchoredInterval{P, T, L, R} where {L<:Intervals.Bounded, R<:Intervals.Bounded}) where {P, T} at /Users/smoritz/.julia/packages/Intervals/ua9cq/src/anchoredinterval.jl:181
  convert(::Type{T}, ::Factorization) where T<:AbstractArray at /Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/factorization.jl:58
  ...
Stacktrace:
  [1] setproperty!(x::OrdinaryDiffEq.ODEIntegrator{Euler, false, ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Nothing, Float64, Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}, Float64, Float64, Float64, Vector{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}}, ODESolution{Float64, 2, Vector{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}}}, ODEProblem{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Tuple{Float64, Float64}, false, Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}, ODEFunction{false, typeof(MitosisStochasticDiffEq.CovarianceFilterODE), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Euler, OrdinaryDiffEq.InterpolationData{ODEFunction{false, typeof(MitosisStochasticDiffEq.CovarianceFilterODE), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Vector{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}}, Vector{Float64}, Vector{Vector{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}}}, OrdinaryDiffEq.EulerConstantCache}, DiffEqBase.DEStats}, ODEFunction{false, typeof(MitosisStochasticDiffEq.CovarianceFilterODE), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, OrdinaryDiffEq.EulerConstantCache, OrdinaryDiffEq.DEOptions{Float64, Float64, Float64, Float64, typeof(DiffEqBase.ODE_DEFAULT_NORM), typeof(opnorm), Nothing, CallbackSet{Tuple{}, Tuple{}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryMinHeap{Float64}, DataStructures.BinaryMinHeap{Float64}, Nothing, Nothing, Int64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Tuple{}, Tuple{}}, ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Float64, Nothing, OrdinaryDiffEq.DefaultInit}, f::Symbol, v::ArrayPartition{Float64, Tuple{SVector{2, Float64}, Matrix{Float64}, Vector{Float64}}})
    @ Base ./Base.jl:34
  [2] initialize!(integrator::OrdinaryDiffEq.ODEIntegrator{Euler, false, ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Nothing, Float64, Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}, Float64, Float64, Float64, Vector{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}}, ODESolution{Float64, 2, Vector{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}}}, ODEProblem{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Tuple{Float64, Float64}, false, Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}, ODEFunction{false, typeof(MitosisStochasticDiffEq.CovarianceFilterODE), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Euler, OrdinaryDiffEq.InterpolationData{ODEFunction{false, typeof(MitosisStochasticDiffEq.CovarianceFilterODE), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Vector{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}}, Vector{Float64}, Vector{Vector{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}}}, OrdinaryDiffEq.EulerConstantCache}, DiffEqBase.DEStats}, ODEFunction{false, typeof(MitosisStochasticDiffEq.CovarianceFilterODE), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, OrdinaryDiffEq.EulerConstantCache, OrdinaryDiffEq.DEOptions{Float64, Float64, Float64, Float64, typeof(DiffEqBase.ODE_DEFAULT_NORM), typeof(opnorm), Nothing, CallbackSet{Tuple{}, Tuple{}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryMinHeap{Float64}, DataStructures.BinaryMinHeap{Float64}, Nothing, Nothing, Int64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Tuple{}, Tuple{}}, ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Float64, Nothing, OrdinaryDiffEq.DefaultInit}, cache::OrdinaryDiffEq.EulerConstantCache)
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/5egkj/src/perform_step/fixed_timestep_perform_step.jl:47
  [3] __init(prob::ODEProblem{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Tuple{Float64, Float64}, false, Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}, ODEFunction{false, typeof(MitosisStochasticDiffEq.CovarianceFilterODE), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, alg::Euler, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{Val{true}}; saveat::Tuple{}, tstops::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Nothing, callback::Nothing, dense::Bool, calck::Bool, dt::Float64, dtmin::Nothing, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Int64, abstol::Float64, reltol::Float64, qmin::Int64, qmax::Int64, qsteady_min::Int64, qsteady_max::Int64, qoldinit::Int64, fullnormalize::Bool, failfactor::Int64, beta1::Nothing, beta2::Nothing, maxiters::Int64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), internalopnorm::typeof(opnorm), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias_u0::Bool, alias_du0::Bool, initializealg::OrdinaryDiffEq.DefaultInit, kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/5egkj/src/solve.jl:433
  [4] #__solve#404
    @ ~/.julia/packages/OrdinaryDiffEq/5egkj/src/solve.jl:4 [inlined]
  [5] #solve_call#56
    @ ~/.julia/packages/DiffEqBase/qntkj/src/solve.jl:61 [inlined]
  [6] solve_up(prob::ODEProblem{ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, Tuple{Float64, Float64}, false, Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}, ODEFunction{false, typeof(MitosisStochasticDiffEq.CovarianceFilterODE), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, sensealg::Nothing, u0::ArrayPartition{Float64, Tuple{Vector{Float64}, Matrix{Float64}, Vector{Float64}}}, p::Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}, args::Euler; kwargs::Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:tstops, :abstol, :reltol), Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Float64, Float64}}})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/qntkj/src/solve.jl:82
  [7] #solve#57
    @ ~/.julia/packages/DiffEqBase/qntkj/src/solve.jl:70 [inlined]
  [8] _backwardfilter(filter::MitosisStochasticDiffEq.CovarianceFilter, k::MitosisStochasticDiffEq.SDEKernel{AffineMap{Matrix{Float64}, SVector{2, Float64}}, ConstantMap{Diagonal{Float64, SVector{2, Float64}}}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}}, ::Tuple{Float64, Vector{Float64}, Matrix{Float64}}; alg::Euler, inplace::Bool, apply_timechange::Bool, abstol::Float64, reltol::Float64)
    @ MitosisStochasticDiffEq ~/.julia/dev/MitosisStochasticDiffEq/src/filter.jl:94
  [9] backwardfilter(k::MitosisStochasticDiffEq.SDEKernel{AffineMap{Matrix{Float64}, SVector{2, Float64}}, ConstantMap{Diagonal{Float64, SVector{2, Float64}}}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}}, Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}}, p::WGaussian{(:μ, :Σ, :c), Tuple{Vector{Float64}, Matrix{Float64}, Float64}}; filter::MitosisStochasticDiffEq.CovarianceFilter, alg::Euler, inplace::Bool, apply_timechange::Bool, abstol::Float64, reltol::Float64)
    @ MitosisStochasticDiffEq ~/.julia/dev/MitosisStochasticDiffEq/src/filter.jl:12
 [10] bwfiltertree!(Q::Vector{Union{Missing, WGaussian{(:μ, :Σ, :c), Tuple{Vector{Float64}, Matrix{Float64}, Float64}}}}, tree::Tree{Node{UInt16, NewickData{Float64, String}}, Float64, Int64, UInt16, String}, θlin::Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}, dt0::Float64; apply_time_change::Bool)
    @ Main ~/.julia/dev/MitosisStochasticDiffEq/example/phylo3/sdetree.jl:93
 [11] bwfiltertree!(Q::Vector{Union{Missing, WGaussian{(:μ, :Σ, :c), Tuple{Vector{Float64}, Matrix{Float64}, Float64}}}}, tree::Tree{Node{UInt16, NewickData{Float64, String}}, Float64, Int64, UInt16, String}, θlin::Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}, dt0::Float64)
    @ Main ~/.julia/dev/MitosisStochasticDiffEq/example/phylo3/sdetree.jl:82
 [12] mcmc(tree::Tree{Node{UInt16, NewickData{Float64, String}}, Float64, Int64, UInt16, String}, Xd::Vector{SVector{2, Float64}}, f::typeof(f), g::typeof(g), θlin::Tuple{Matrix{Float64}, SVector{2, Float64}, Diagonal{Float64, SVector{2, Float64}}}, θinit::Tuple{Vector{Float64}, SVector{2, Float64}}; ρ::Float64, iters::Int64, dt::Float64, σprop::Float64, precisionatleaves::Float64)
    @ Main ~/.julia/dev/MitosisStochasticDiffEq/example/phylo3/phylo_autobffg3.jl:55
 [13] top-level scope
    @ ~/.julia/dev/MitosisStochasticDiffEq/example/phylo3/phylo_autobffg3.jl:101

Infer diffusivity

Like the drift is conjugate in the linear parameters, the diffusion coefficient is conjugate Wishart in its scale... I have a sketch here https://github.com/mschauer/Mitosis.jl/blob/master/example/wiener.jl