[WIP] ExpectationProblem interface

[agerlach]

This WIP PR is meant to start discussing interface updates and a general restructuring of DEU. NOTE: the interface described below is implemented w/ tests. However, it is not at feature parity. Mainly, I have not defined the integrand yet to support batch evaluation nor performed extensive testing for AD.

Big shout out to @aml5600 for his help.

This addresses, #48, #44, #34,

The Problems

The interface of the current release for expectation requires passing two arrays for the initial conditions and parameters that mix number and distribution types, e.g.

1st, one defines the system

function pend!(du, u, p, t)
    du[1] = u[2]
    du[2] = -p[1]/p[2]*sin(u[1]) + mean(@view p[3:end])
    nothing
end

g(soln) = soln[1, end]
u0 = [0.0, 0.1]
ps = [9.807; 1.0; 0.0] 
tspan = (0.0,10.0)
prob = ODEProblem{true}(pend!, u0, tspan, ps)

then defines, a new array of deterministic and uncertain values for the initial conditions and parameters

u0_dist = [0.0, Uniform(-0.1, 0.1]
ps_dist = [Normal(9.807, .01), Uniform(.9, 1.1), 0.0]
expectation(g, prob, u0_dist, ps_dist, Koopman(), Tsit5())

Problem 1: This is not type stable due to requiring Array{Any,1} for the distributions.
Problem 2 Internally the remake call forces the initial condition and parameters to be Array. This fails if the ODE is written for something like ComponentArrays. Here, all the information for u0 and ps in the DEproblem is lost.
Problem 3 Internally this assumes all distributions are independent so the joint can be formed as the product of pdfs

Proposed New Interface

We now have problem/solve workflow like the rest of SciML by introducing the ExpectationProblem. This essentially stores all the functions and the distribution required to solve problems of the form: ∫ g(S(h(x,u0,p)))*pdf(d,x)dx.

Take 𝕏, 𝕌, ℙ, as the uncertainty, state, and parameters spaces, respectively:

• pdf(d,.): 𝕏 → ℝ ,d, is any object that supports pdf and extrema
• h: 𝕏 × 𝕌 × ℙ → 𝕌 × ℙ ,h is a callable object that maps points from the uncertainty/IC/parameter spaces to a new IC/parameter
• S: 𝕌 × ℙ → 𝕌 ,S is a callable object that maps IC/param to new state(s)
• g: 𝕌 × ℙ → ℝⁿᵒᵘᵗ ,observable

To define, S, we introduce SystemMap as a callable object that wraps remake -> solve, e.g.

prob = ODEProblem(eom, u0, tspan, p)
sm = SystemMap(prob, Tsit5(), callback = cb, save_everystep=false, save_start = false)
sm(u0, p)

is equivalent to

prob = ODEProblem(eom, u0, tspan, p)
soln = solve(prob, Tsit5(), callback = cb, save_everystep=false, save_start = false)

Doing this has 3 primary benefits 1) We can now support arbitrary callable objects for the system map and 2) It prevents downstream args.../kwargs... collisions, by isolating the DE solve arguments in the SystemMap 3) it ensures proper conversion to the correct u0, p types defined by the ODEProblem.

Example 1

using OrdinaryDiffEq, BenchmarkTools, Distributions,
         ComponentArrays, Parameters, DiffEqUncertainty

# 2D projectile w/ nonlinear drag, using ComponentArrays (previously not supported)
function projectile!(du, u, p, t)
    @unpack CdS, m, ρ, g = p
    @unpack x, z, ẋ, ż = u
    coeff = -0.5 * CdS / m * ρ* sqrt(ẋ^2+ż^2)
    du.x = ẋ
    du.z = ż
    du.ẋ = coeff * ẋ
    du.ż = coeff * ż - g
    nothing
end

# Observable, impact x position and terminal velocity
g(soln, p) = [soln[end].x, sqrt(soln[end].ẋ^2 + soln[end].ż^2)] 

## ODE Problem Setup
u0 = ComponentArray(x=0.0, z=0.0, ẋ = 10.0, ż = 10.0)
ps = ComponentArray(CdS = 1.0, m = 2.0, ρ = 1.22, g = 9.807)
tspan = (0.0,1000.0)

ground_impact(u,t,integrator) = u.z
affect!(integrator) = terminate!(integrator)
cb = ContinuousCallback(ground_impact,affect!);

prob = ODEProblem(projectile!, u0, tspan, ps)
sm = SystemMap(prob, Tsit5(); callback = cb, save_everystep=false, save_start = false)

# Define Distribution, can now handle any univariate/MV distribution! 
# defined as initial velocity, elevation, CdS
dist = Product([Uniform(13.0, 15.0), Uniform(.9*π/4, 1.1*π/4), Truncated(Normal(1.0, .02), .9, 1.1)])

# Define function of how to change state/param given point in uncertainty space, note how change of variable required from uncertainty space to state space!
function h(q, u0, p)
    # q = [velocity, elevation, CdS]
    u0.ẋ = q[1]*cos(q[2])
    u0.ż = q[1]*cos(q[2])
    p.CdS = q[3]
    u0, p
end

exprob = ExpectationProblem(sm, g, h, dist; nout = 2)
solve(exprob, Koopman())
solve(exprob, MonteCarlo(10000))

Example 2

Because the ExpectationProblem just operates on callable objects we can also solve for expectation of arbitrary functions

using OrdinaryDiffEq, BenchmarkTools, Distributions, DiffEqUncertainty, LinearAlgebra

g(x, p) = dot(p,sin.(x))
dist = Product(Uniform.(0, ones(3)))
p = [1.0, 2.0, 3.0]

exprob = ExpectationProblem(g, dist, p)
solve(exprob, Koopman())
solve(exprob, MonteCarlo(1000))

Arbitrary Distribution Support

What if a user has some complicated distribution they want to consider that is not supported by Distributions.jl? For that we introduce the GenericDistribution that implements the required Distributions.jl interface for usage in DiffEqUncertainty. The user just supplies a function for computing the pdf, a rand() function (if using MonteCarlo), and lower/upper bounds.

A trivial example looks like

d1 = MvNormal([0,0], [1 3/5; 3/5 2])
d2 = Uniform(1,2)
pdf_func(x) = pdf(d1, @view x[1:2]) * pdf(d2, x[3])
rand_func() = [rand(d1); rand(d2)]
lb = [minimum(d1); minimum(d2)]
ub = [maximum(d1); maximum(d2)]
gd = GenericDistribution(pdf_func, rand_func, lb, ub)
exprob = ExpectationProblem(..., gd)

Non-Quadrature.jl integration

exprob = ExpectationProblem(...)
f = build_integrand(exprob, Koopman())

Will return the integrand required to solve the Koopman Expectation problem. The user can then use non-Quadrature.jl integration methods.

Other Changes of Note

1. The code no longer forces the initial condition and parameters to be of type Array. They now will match the types of prob.u0 and prob.p. This enables its usage for system written using <:AbstractArrays
2. Everything is type stable except for upstream issues w/ Quadrature.jl
3. Inferrence tests are included
4. I have several Zygote adjoints implemented to test several options independent of Quadrature.jl. These include the options to fuse the primal in the gradient integrand and to only consider the norm with respect to the primal, which should more closely match the evaluation points of doing ForwardDiff. These are here for testing and would get rolled into Quadrature.jl when ready.
5. These Zygote adjoints are not type stable, but I think that is because I am hitting improve inference for getproperty FluxML/Zygote.jl#909 (comment)

To Do

• Batch Support
  • DiffEq Problems
  • General Functions
• Define common solve output type for MonteCarlo() and Koopman(), possibly store ODESolutions
• Define CentralMomentProblem
• Update Readme
• Update SciMLTutorials
• Add Doc strings
• Test AD

[agerlach]

@ChrisRackauckas I would like to get your thoughts on this before pushing forward.

[ChrisRackauckas]

I think this is good. I don't quite see the need for SystemMap as it's equivalent to (u0,p) -> ... with solve, which I think is extra machinery that isn't quite needed. Also, with this change I think we could instead make this be a whole Expectations.jl, with differential equation examples, but not necessarily tied to differential equations.

[agerlach]

@aml5600 and I debated on whether or not to include a non-DE interface. I wanted to have it for a couple of trivial applications, but I can see not wanting to maintain that. Alternatively, if we only had a DE interface, on could simply get around it by simulating from from t = 0 to 0.

RE: SystemMap I don't disagree. Would you prefer to just have the user pass the (u0,p)-> closure or just build it on construction? The advantage I see of using SystemMap or the closure is an isolation of the kwargs that prevent collisions, e.g. reltol, abstol for DE solve vs quadrature.

[ChrisRackauckas]

RE: SystemMap I don't disagree. Would you prefer to just have the user pass the (u0,p)-> closure or just build it on construction? The advantage I see of using SystemMap or the closure is an isolation of the kwargs that prevent collisions, e.g. reltol, abstol for DE solve vs quadrature.

I think the issue is always that it confuses people when you have a non-limitation which becomes a documentation limitation because of practice. SystemMap in documentation I think is not much different from (u0,p)->, but if the tutorials use SystemMap, 99% of users will think it only works with a SystemMap input. Since it's not doing anything crucial, I'd just remove it from the equation.

I debated on whether or not to include a non-DE interface. I wanted to have it for a couple of trivial applications, but I can see not wanting to maintain that. Alternatively, if we only had a DE interface, on could simply get around it by simulating from from t = 0 to 0.

If it's just functions of (u0,p) then it's a bit more general, and the user just encloses tspan.

[jmurphy6895]

@ChrisRackauckas I've been discussing with @agerlach about getting involved in the development of this and I wanted to put on your radar that I'm available to help build up tests or address other issues as this pull request and library move forward

[ChrisRackauckas]

what happened to the probints one?

[agerlach]

These tests took a long time to run locally, so I temp commented them out.

[ChrisRackauckas]

I dug through this, and I think we can just merge it if it gets rebased and we figure out why the probints tests were commented out. Of course more can keep being done, but I think this is a big enough improvement to just go for it.

[agerlach]

This is not a feature parity yet, so I'm hesitant to merge as is. Looking of the todo list, it looks like CentralMomentProblem support and AD Tests are really to the two big outstanding parts. I also think we should consider some common output type either as part of this PR or prior to the next release, that way we don't have a major breaking release in the near term

[agerlach]

@jmurphy6895 Adding central moment support via a CentralMomentProblem for both Koopman and Monte Carlo would be a big help

[agerlach]

@ChrisRackauckas We never settled on the interface wrt SystemMap. My primary reason for it was to separate kwargs between the ODE solve and the Quadrature solve. I have zero issue removing, but I do want brainstorm alternatives. With that said it is under SciML and I want to make sure to respects the common workflow.

[baggepinnen]

These lines are technically type piracy right? Also, they might not be correct for any AbstractNormal that accepts singular covariance matrices, e.g., Diagonal([0, 1]) does not have infinite support for the first variable.

[agerlach]

Yes, the UnivariateKDE parts should be upstreamed. Base.minimum(d::AbstractMvNormal) was upstreamed in Distributions in JuliaStats/Distributions.jl#1319, probably need to double check the PR discussion to make sure a bug wasn't introduced.

[baggepinnen]

These are already defined in Distributions
https://github.com/JuliaStats/Distributions.jl/blob/master/src/multivariate/product.jl#L43

[agerlach]

I had these here as a stop gap until this PR when through JuliaStats/Distributions.jl#1319

[baggepinnen]

I'm not sure about the semantics of sol[:], but it looks like there are unnecessary allocations from the [:] here.
Since they are marked @inline, I guess performance matters here?

[agerlach]

set_result! ends up formating the solution of an EnsembleSolve to the format required by Quadrature.jl. This is getting called once per batch integrand evaluation. I'm not sure how to clean up the allocation.

[agerlach]

This also highlights an interface issue with Quadrature (which I am probably guilty of creating), i.e. we shouldn't really need separate set_result! for scalar and vector valued integrands.

[agerlach]

Actions from our call today:

• Remove SystemMap
• Named tuple for DE vs Quad kwargs
• Common output type w/ members u, resid, original
• Uncomment AD tests and mark as broken where needed
• Remove blocked comments

[baggepinnen]

Why does h need to take samples as well as IC and parameters?

h: 𝕏 × 𝕌 × ℙ → 𝕌 × ℙ ,h is a callable object that maps points from the uncertainty/IC/parameter spaces to a new IC/parameter

wouldn't it be enough for h to take x? Is it to be able to modify (u, p) in place? If so, should h be called h!(u0, p, q) in the docs to make this clear?

[agerlach]

Consider the ballistic projectile example in the OP. Here, the dynamics are best expressed in cartesian coordinates consisting of a down range and vertical position and corresponding velocities. If you have uncertainty in elevation angle you need to do a change of variables into the state space. However, this process does not effect the entire state-space as the initial launch location is not uncertain. So, if you discretize the uncertainty space, you need a way to map that to the subset of the state-space that is effected in addition to the subset of the parameter space if needed. Here, I map uncertain elevation and exit velocity to the velocity states while leaving position unchanged.

I'm not sure if you can do this in-place. The ODEProblem is passed by reference in the prob_func here https://github.com/agerlach/DiffEqUncertainty.jl/blob/50024a855d522f3ba468a9cd41419196b633f27e/src/expectation/expectation.jl#L60 used by EnsembleProblem. In-place would modify the ODEProblem prior to calling remake

# Define Distribution, can now handle any univariate/MV distribution! 
# defined as initial velocity, elevation, CdS
dist = Product([Uniform(13.0, 15.0), Uniform(.9*π/4, 1.1*π/4), Truncated(Normal(1.0, .02), .9, 1.1)])

# Define function of how to change state/param given point in uncertainty space
function h(q, u0, p)
    # q = [velocity, elevation, CdS]
    u0.ẋ = q[1]*cos(q[2])
    u0.ż = q[1]*cos(q[2])
    p.CdS = q[3]
    u0, p
end

[ChrisRackauckas]

I played around with removing SystemMap but noticed that indeed it's a PITA so it makes sense to keep it as a helper function. It does do a few bits of nice things that are not so easy to automate otherwise.