Problem Interface

src/DiffEqUncertainty.jl

@@ -2,10 +2,16 @@ module DiffEqUncertainty
 
 using DiffEqBase, Statistics, Distributions, Reexport
 @reexport using Quadrature
+
+abstract type AbstractUncertaintyProblem end
+
+include("uncertainty_utils.jl")
+include("uncertainty_problems.jl")
 include("probints.jl")
 include("koopman.jl")
 
 export ProbIntsUncertainty,AdaptiveProbIntsUncertainty
+export AbstractUncertaintyProblem, ExpectationProblem
 export expectation, centralmoment, Koopman, MonteCarlo
 
 end

src/koopman.jl

@@ -7,13 +7,88 @@ struct MonteCarlo <: AbstractExpectationAlgorithm end
 @inline tuplejoin(x, y) = (x..., y...)
 @inline tuplejoin(x, y, z...) = (x..., tuplejoin(y, z...)...)
 
-_rand(x::T) where T <: Sampleable = rand(x)
-_rand(x) = x
-
 function __make_map(prob::ODEProblem, args...; kwargs...)
     (u,p) -> solve(remake(prob,u0=u,p=p), args...; kwargs...)
 end
 
+"""
+    solve(prob, expalg::Koopman, args; kwargs...)
+
+Solves the ExpectationProblem using via the Koopman expectation
+
+Both args and kwargs are passed to DifferentialEquation solver
+
+Special kwargs:
+    - maxiters: max quadrature iterations (default 1000000)
+    - batch: solve quadrature using n-batch processing (default 0, off)
+    - quadalg: quadrature algorithm to use (default HCubatureJL())
+    - ireltol, iabstol: integration relative and absolute tolerances (default 1e-2, 1e-2)
+"""
+function DiffEqBase.solve(prob::ExpectationProblem, expalg::Koopman, args...;
+    maxiters=1000000,
+    batch=0,
+    quadalg=HCubatureJL(),
+    ireltol=1e-2, iabstol=1e-2,
+    kwargs...)
+
+    jargs = tuplejoin(args)
+    jkwargs = tuplejoin(prob.kwargs, kwargs)
+
+    # build the integrand for ∫(Ug)(x) * f(x) dx 
+    integrand = function(fq, xq, pq)
+        # convert to physical x and p
+        (_x, _p) = prob.comp_func(prob.to_phys(xq,pq)...)
+
+        # solve ODE and evaluate observable
+        if batch == 0
+            # scalar solution
+            _prob = remake(prob.ode_prob, u0=_x, p=_p)
+            Ug = prob.g(solve(_prob, jargs...; jkwargs...))
+            f0 = prob.f0_func(_x, _p)
+        else
+            # ensemble solution
+            prob_func(prob, i, repeat) = remake(prob, u0=_x[:,i], p=_p[:,i])
+            output_func(sol,i) = prob.g(sol), false
+            _prob = EnsembleProblem(prob.ode_prob, prob_func=prob_func, output_func=output_func)
+            Ug = solve(_prob, jargs...; trajectories=size(xq,2), jkwargs...)[:]
+            f0 = map(prob.f0_func, zip(_x,_p))
+        end
+
+        fq .= Ug .* f0
+        return nothing
+    end
+
+    # solve the integral using quadrature methods
+    intprob = QuadratureProblem(integrand, prob.quad_lb, prob.quad_ub, prob.p_quad, batch=batch, nout=prob.nout)
+    sol = solve(intprob, quadalg, reltol=ireltol, abstol=iabstol, maxiters=maxiters)
+end
+
+"""
+    solve(prob, expalg::MonteCarlo, args; kwargs...)
+
+Solves the ExpectationProblem using via the Monte Carlo integration
+
+Both args and kwargs are passed to DifferentialEquation solver
+
+"""
+function DiffEqBase.solve(prob::ExpectationProblem, expalg::MonteCarlo, args...; trajectories,kwargs...)
+    jargs = tuplejoin(prob.args, args)
+    jkwargs = tuplejoin(prob.kwargs, kwargs)
+
+    prob_func = function (prob, i, repeat)
+        _u0, _p = prob.comp_func(prob.samp_func()...)
+        remake(prob, u0=_u0, p=_p)
+    end
+
+    output_func(sol, i) = (prob.g(sol), false)
+
+    monte_prob = EnsembleProblem(prob.ode_prob;
+                output_func=output_func,
+                prob_func=prob_func)
+    sol = solve(monte_prob, jargs...;trajectories=trajectories,jkwargs...)
+    mean(sol.u)
+end
+
 function expectation(g::Function, prob::ODEProblem, u0, p, expalg::Koopman, args...;
                         u0_CoV=(u,p)->u, p_CoV=(u,p)->p,
                         maxiters=1000000,

src/uncertainty_problems.jl

@@ -0,0 +1,77 @@
+
+struct ExpectationProblem{T, Tg, Tq, Tp, Tf, Ts, Tc, Tb, Tr, To, Tk} <: AbstractUncertaintyProblem
+    Tscalar::Type{T}
+    nout::Int64
+    g::Tg
+    to_quad::Tq
+    to_phys::Tp
+    f0_func::Tf
+    samp_func::Ts
+    comp_func::Tc
+    quad_lb::Tb
+    quad_ub::Tb
+    p_quad::Tr
+    ode_prob::To
+    kwargs::Tk
+end
+
+DEFAULT_COMP_FUNC(x,p) = (x,p)
+
+# Builds problem from (arrays) of u0 and p distribution(s)
+function ExpectationProblem(g::Function, u0_dist, p_dist, prob::ODEProblem, nout=1; 
+    comp_func=DEFAULT_COMP_FUNC, lower_bounds=nothing, upper_bounds=nothing, kwargs...)
+
+    _xdist = u0_dist isa AbstractArray ? u0_dist : [u0_dist]
+    _pdist = p_dist isa AbstractArray ? p_dist : [p_dist] 
+
+    T = promote_type(eltype.(mean.([_xdist...,_pdist...]))...)
+
+    # build shuffle/unshuffle functions
+    usizes = Vector{Int64}([length(u) for u in _xdist])
+    psizes = Vector{Int64}([length(p) for p in _pdist])
+    nu = sum(usizes)
+    np = sum(psizes)
+    mx = vcat([repeat([u isa Distribution],length(u)) for u in _xdist]...)
+    mp = vcat([repeat([p isa Distribution],length(p)) for p in _pdist]...)
+    dist_mask = [mx..., mp...]
+
+    # map physical x-state, p-params to quadrature state and params
+    to_quad = function(x,p)
+        esv = [x;p]
+        return (view(esv,dist_mask), view(esv, .!(dist_mask)))
+    end
+
+    # map quadrature x-state, p-params to physical state and params
+    to_phys = function(x,p)
+        x_it, p_it = 0, 0
+        esv = map(1:length(dist_mask)) do idx
+            dist_mask[idx] ? T(x[x_it+=1]) : T(p[p_it+=1])
+        end
+        return (view(esv,1:nu), view(esv,(nu+1):(nu+np)))
+    end
+
+    # evaluate the f0 (joint) distribution
+    f0_func = function(u,p)
+        fu = prod([_pdf(dist, u[idx]) for (idx,dist) in zip(accumulated_range(usizes), _xdist)])
+        fp = prod([_pdf(dist, p[idx]) for (idx,dist) in zip(accumulated_range(psizes), _pdist)])
+        return fu * fp
+    end
+
+    # sample from (joint) distribution
+    samp_func() = comp_func(vcat(_rand.(_xdist)...), vcat(_rand.(_pdist)...))
+
+    # compute the bounds
+    lb = isnothing(lower_bounds) ? to_quad(comp_func(vcat(_minimum.(_xdist)...), vcat(_minimum.(_pdist)...))...)[1] : lower_bounds
+    ub = isnothing(upper_bounds) ? to_quad(comp_func(vcat(_maximum.(_xdist)...), vcat(_maximum.(_pdist)...))...)[1] : upper_bounds
+
+    # compute "static" quadrature parameters
+    p_quad = to_quad(comp_func(vcat(mean.(_xdist)...), vcat(mean.(_pdist)...))...)[2]
+
+
+    return ExpectationProblem(T,nout,g,to_quad,to_phys,f0_func,samp_func,comp_func,lb,ub,p_quad,prob,kwargs)
+end
+
+# Builds problem from (array) of u0 distribution(s)
+function ExpectationProblem(g::Function, u0_dist, prob::ODEProblem, nout=1; kwargs...)
+    return ExpectationProblem(g,u0_dist,[],prob,nout=nout; kwargs...)
+end

src/uncertainty_utils.jl

@@ -0,0 +1,19 @@
+
+# wraps Distribtuions.jl
+# note: MultivariateDistributions do not support minimum/maximum. Setting to -Inf/Inf with
+# the knowledge that it will cause quadrature integration to fail if Koopman is used. To use
+# MultivariateDistributions the upper/lower bounds should be set with the kwargs.
+_minimum(f::T) where T <: MultivariateDistribution = -Inf .* ones(eltype(f), size(f)...)
+_minimum(f) = minimum(f)
+_maximum(f::T) where T <: MultivariateDistribution = Inf .* ones(eltype(f), size(f)...)
+_maximum(f) = maximum(f)
+_rand(f::T) where T <: Distribution = rand(f)
+_rand(x) = x
+_pdf(f::T, x) where T <: Distribution = pdf(f,x)
+_pdf(f, x) = one(eltype(x))
+
+# creates a tuple of idices, or ranges, from array partition lengths
+function accumulated_range(partition_lengths)
+    c = [0, cumsum(partition_lengths)...]
+    return Tuple(c[i]+1==c[i+1] ? c[i+1] : (c[i]+1):c[i+1] for i in 1:length(c)-1)
+end

test/koopman.jl

@@ -21,6 +21,17 @@ prob = ODEProblem(eom!,u0,tspan,p)
 A = [0.0 1.0; -p[1] -p[2]]
 u0s_dist = [Uniform(1,10), Uniform(2,6)]
 
+@testset "Koopman solve, nout = 1" begin
+  g(sol) = sol[1,end]
+  exp_prob = ExpectationProblem(g, u0s_dist, p, prob)
+  analytical = (exp(A*tspan[end])*mean.(u0s_dist))[1]
+
+  for alg ∈ quadalgs
+    @info "$alg"
+    @test solve(exp_prob, Koopman(), Tsit5(); quadalg=alg)[1] ≈ analytical rtol=1e-2
+  end
+end
+
 @testset "Koopman Expectation, nout = 1" begin
   g(sol) = sol[1,end]
   analytical = (exp(A*tspan[end])*mean.(u0s_dist))[1]

test/problems.jl

@@ -0,0 +1,85 @@
+using DiffEqUncertainty
+using OrdinaryDiffEq
+using Distributions
+
+
+function eom!(du,u,p,t)
+  @inbounds begin
+    du[1] = u[2]
+    du[2] = -p[1]*u[1] - p[2]*u[2]
+  end
+  nothing
+end
+
+u0 = [1.0;1.0]
+tspan = (0.0,3.0)
+p = [1.0, 2.0]
+ode_prob = ODEProblem(eom!,u0,tspan,p)
+g(sol) = sol[1,end]
+
+# various u0/p distributions
+u0_const = [5., 5.]
+u0_mixed = [5., Uniform(1.,9.)]
+u0_mv_dist = MvNormal([5., 5.], [1., 1.])
+p_const = [2., 1f0]
+p_mixed = [Uniform(1.,2.), 1f0]
+
+@testset "Expectation problems" begin
+    @testset "mixed u0, const p" begin 
+        prob = ExpectationProblem(g, u0_mixed, p_const, ode_prob)
+        @test prob.Tscalar == Float64
+        @test prob.to_quad([1., 2.], [3., 4.]) == ([2.], [1.,3.,4.])
+        @test prob.to_phys([2.], [1.,3.,4.]) == ([1., 2.], [3., 4.])
+        @test prob.f0_func(u0_const, p_const) == (1. / 8.)
+        (x,p) = prob.samp_func()
+        @test (x[1] == 5.) && (1. <= x[2] <= 9.) && (p[1] == 2.) && (p[2] == 1.)
+        @test (prob.quad_lb, prob.quad_ub) == ([1.], [9.])
+        @test prob.p_quad == [5., 2., 1.]
+    end
+
+    @testset "mixed u0, mixed p" begin
+        prob = ExpectationProblem(g, u0_mixed, p_mixed, ode_prob)
+        @test prob.Tscalar == Float64
+        @test prob.to_quad([1., 2.], [3., 4.]) == ([2.,3.], [1., 4.])
+        @test prob.to_phys([2., 3.], [1., 4.]) == ([1., 2.], [3., 4.])
+        @test prob.f0_func(u0_const, p_const) == (1. / 8.)
+        (x,p) = prob.samp_func()
+        @test (x[1] == 5.) && (1. <= x[2] <= 9.) && (1. <= p[1] <= 2.) && (p[2] == 1.)
+        @test (prob.quad_lb, prob.quad_ub) == ([1.,1.,], [9.,2.])
+        @test prob.p_quad == [5., 1.]
+    end
+
+    @testset "multivariate u0, mixed p" begin
+        prob = ExpectationProblem(g, u0_mv_dist, p_mixed, ode_prob)
+        @test prob.Tscalar == Float64
+        @test prob.to_quad([1., 2.], [3., 4.]) == ([1.,2.,3.], [4.])
+        @test prob.to_phys([1.,2.,3.], [4.]) == ([1., 2.], [3., 4.])
+        @test prob.f0_func(u0_const, p_const) == pdf(u0_mv_dist, u0_const)
+        @test length.(prob.samp_func()) == (2, 2)
+        @test (prob.quad_lb, prob.quad_ub) == ([-Inf,-Inf,1.], [Inf, Inf,2.])
+        @test prob.p_quad == [1.]
+    end
+
+    @testset "multivariate u0, mixed p w/ bounds" begin
+        prob = ExpectationProblem(g, u0_mv_dist, p_mixed, ode_prob; 
+                lower_bounds=[1.,1.,1.], upper_bounds=[9.,9.,2.])
+        @test (prob.quad_lb, prob.quad_ub) == ([1.,1.,1.], [9.,9.,2.])
+    end
+
+    @testset "completeion function" begin
+        comp_func(x,p) = (x=[2*x[2];x[2]], p=[p[1],2*p[1]])
+        prob = ExpectationProblem(g, u0_mixed, p_mixed, ode_prob; comp_func=comp_func)
+        @test prob.f0_func(u0_const, p_const) == (1. / 8.)
+        (x,p) = prob.samp_func()
+        @test (x[1] == 2*x[2]) && (1. <= x[2] <= 9.) && (1. <= p[1] <= 2.) && (p[2] == 2*p[1])
+        @test (prob.quad_lb, prob.quad_ub) == ([1.,1.], [9.,2.])
+    end
+
+    @testset "no parameters" begin
+        F(x,p,t) = -x
+        ode_nop = ODEProblem(F, u0, tspan)
+        prob = ExpectationProblem(g, u0_mixed, ode_nop)
+        @test prob.to_quad([1.,2.],[]) == ([2.], [1.])
+        @test prob.to_phys([2.],[1.]) == ([1.,2.],[])
+    end
+end

test/runtests.jl

@@ -1,4 +1,5 @@
 using Test
 
 @time @testset "ProbInts tests" begin include("probints.jl") end
+@time @testset "Problem Tests" begin include("problems.jl") end
 @time @testset "Koopman Tests" begin include("koopman.jl") end