[WIP] ExpectationProblem interface

.JuliaFormatter.toml

@@ -0,0 +1 @@
+style = "sciml"

.gitignore

@@ -2,3 +2,4 @@
 *.jl.*.cov
 *.jl.mem
 Manifest.toml
+tests/Manifest.toml

Project.toml

@@ -6,31 +6,21 @@ version = "1.9.0"
 [deps]
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b"
+Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 DiffEqBase = "6"
-Distributions = "0.18, 0.19, 0.20, 0.21, 0.22, 0.23, 0.24, 0.25"
+Distributions = "0.23, 0.24, 0.25"
 KernelDensity = "0.6"
 Integrals = "3.1"
 Reexport = "0.2, 1.0"
 julia = "1.6"
-
-[extras]
-DiffEqGPU = "071ae1c0-96b5-11e9-1965-c90190d839ea"
-DiffEqProblemLibrary = "a077e3f3-b75c-5d7f-a0c6-6bc4c8ec64a9"
-DiffEqSensitivity = "41bf760c-e81c-5289-8e54-58b1f1f8abe2"
-FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-IntegralsCuba = "e00cd5f1-6337-4131-8b37-28b2fe4cd6cb"
-IntegralsCubature = "c31f79ba-6e32-46d4-a52f-182a8ac42a54"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-
-[targets]
-test = ["DiffEqProblemLibrary", "Test", "OrdinaryDiffEq", "IntegralsCuba", "IntegralsCubature", "FiniteDiff", "ForwardDiff", "Zygote", "DiffEqGPU", "DiffEqSensitivity", "LinearAlgebra"]

src/DiffEqUncertainty.jl

@@ -1,18 +1,38 @@
 module DiffEqUncertainty
 
-using DiffEqBase, Statistics, Distributions, Reexport
+# LinearAlgebra
+using DiffEqBase, Statistics, Reexport, RecursiveArrayTools, StaticArrays,
+      Distributions, KernelDensity, Zygote, LinearAlgebra, Random
+using Parameters: @unpack
+
 @reexport using Integrals
-using KernelDensity
+import DiffEqBase: solve
+
+include("expectation/system_utils.jl")
+include("expectation/distribution_utils.jl")
+include("expectation/problem_types.jl")
+include("expectation/solution_types.jl")
+include("expectation/expectation.jl")
 
 include("probints.jl")
-include("koopman.jl")
 
 # Type Piracy, should upstream
-Base.eltype(K::UnivariateKDE)  = eltype(K.density)
+Base.eltype(K::UnivariateKDE) = eltype(K.density)
 Base.minimum(K::UnivariateKDE) = minimum(K.x)
 Base.maximum(K::UnivariateKDE) = maximum(K.x)
+Base.extrema(K::UnivariateKDE) = minimum(K), maximum(K)
+
+Base.minimum(d::AbstractMvNormal) = fill(-Inf, length(d))
+Base.maximum(d::AbstractMvNormal) = fill(Inf, length(d))
+Base.extrema(d::AbstractMvNormal) = minimum(d), maximum(d)
+
+Base.minimum(d::Product) = minimum.(d.v)
+Base.maximum(d::Product) = maximum.(d.v)
+Base.extrema(d::Product) = minimum(d), maximum(d)
+
+export ProbIntsUncertainty, AdaptiveProbIntsUncertainty
 
-export ProbIntsUncertainty,AdaptiveProbIntsUncertainty
-export expectation, centralmoment, Koopman, MonteCarlo
+export Koopman, MonteCarlo, PrefusedAD, PostfusedAD, NonfusedAD
+export GenericDistribution, SystemMap, ExpectationProblem, build_integrand
 
 end

src/expectation/distribution_utils.jl

@@ -0,0 +1,29 @@
+"""
+`GenericDistribution(d, ds...)`
+
+Defines a generic distribution that just wraps functions for pdf function, rand and bounds.
+User can use this for define any arbitray joint pdf. Included b/c Distributions.jl Product
+method of mixed distirbutions are type instable
+"""
+struct GenericDistribution{TF, TRF, TLB, TUB}
+    pdf_func::TF
+    rand_func::TRF
+    lb::TLB
+    ub::TUB
+end
+
+function GenericDistribution(d, ds...)
+    dists = (d, ds...)
+    pdf_func(x) = exp(sum(logpdf(f, y) for (f, y) in zip(dists, x)))
+    rand_func() = [rand(d) for d in dists]
+    lb = SVector(map(minimum, dists)...)
+    ub = SVector(map(maximum, dists)...)
+
+    GenericDistribution(pdf_func, rand_func, lb, ub)
+end
+
+Distributions.pdf(d::GenericDistribution, x) = d.pdf_func(x)
+Base.minimum(d::GenericDistribution) = d.lb
+Base.maximum(d::GenericDistribution) = d.ub
+Base.extrema(d::GenericDistribution) = minimum(d), maximum(d)
+Base.rand(d::GenericDistribution) = d.rand_func()

src/expectation/expectation.jl

@@ -0,0 +1,214 @@
+abstract type AbstractExpectationADAlgorithm end
+struct NonfusedAD <: AbstractExpectationADAlgorithm end
+struct PrefusedAD <: AbstractExpectationADAlgorithm
+    norm_partials::Bool
+end
+PrefusedAD() = PrefusedAD(true)
+struct PostfusedAD <: AbstractExpectationADAlgorithm
+    norm_partials::Bool
+end
+PostfusedAD() = PostfusedAD(true)
+
+abstract type AbstractExpectationAlgorithm <: DiffEqBase.DEAlgorithm end
+struct Koopman{TS} <:
+       AbstractExpectationAlgorithm where {TS <: AbstractExpectationADAlgorithm}
+    sensealg::TS
+end
+Koopman() = Koopman(NonfusedAD())
+struct MonteCarlo <: AbstractExpectationAlgorithm
+    trajectories::Int
+end
+
+# Builds integrand for arbitrary functions
+function build_integrand(prob::ExpectationProblem, ::Koopman, ::Val{false})
+    @unpack g, d = prob
+    function (x, p)
+        g(x, p) * pdf(d, x)
+    end
+end
+
+# Builds integrand for DEProblems
+function build_integrand(prob::ExpectationProblem{F}, ::Koopman,
+                         ::Val{false}) where {F <: SystemMap}
+    @unpack S, g, h, d = prob
+    function (x, p)
+        ū, p̄ = h(x, p.x[1], p.x[2])
+        g(S(ū, p̄), p̄) * pdf(d, x)
+    end
+end
+
+function _make_view(x::Union{Vector{T}, Adjoint{T, Vector{T}}}, i) where {T}
+    @view x[i]
+end
+
+function _make_view(x, i)
+    @view x[:, i]
+end
+
+function build_integrand(prob::ExpectationProblem{F}, ::Koopman,
+                         ::Val{true}) where {F <: SystemMap}
+    @unpack S, g, h, d = prob
+
+    if prob.nout == 1 # TODO fix upstream in quadrature, expected sizes depend on quadrature method is requires different copying based on nout > 1
+        set_result! = @inline function (dx, sol)
+            dx[:] .= sol[:]
+        end
+    else
+        set_result! = @inline function (dx, sol)
+            dx .= reshape(sol[:, :], size(dx))
+        end
+    end
+
+    prob_func = function (prob, i, repeat, x)  # TODO is it better to make prob/output funcs outside of integrand, then call w/ closure?
+        u0, p = h((_make_view(x, i)), prob.u0, prob.p)
+        remake(prob, u0 = u0, p = p)
+    end
+
+    output_func(sol, i, x) = (g(sol, sol.prob.p) * pdf(d, (_make_view(x, i))), false)
+
+    function (dx, x, p) where {T}
+        trajectories = size(x, 2)
+        # TODO How to inject ensemble method in solve? currently in SystemMap, but does that make sense?
+        ensprob = EnsembleProblem(S.prob; output_func = (sol, i) -> output_func(sol, i, x),
+                                  prob_func = (prob, i, repeat) -> prob_func(prob, i,
+                                                                             repeat, x))
+        sol = solve(ensprob, S.args...; trajectories = trajectories, S.kwargs...)
+        set_result!(dx, sol)
+        nothing
+    end
+end
+
+# solve expectation problem of generic callable functions via MonteCarlo
+function DiffEqBase.solve(exprob::ExpectationProblem, expalg::MonteCarlo)
+    params = parameters(exprob)
+    dist = distribution(exprob)
+    g = observable(exprob)
+    ExpectationSolution(mean(g(rand(dist), params) for _ in 1:(expalg.trajectories)), nothing, nothing)
+end
+
+# solve expectation over DEProblem via MonteCarlo
+function DiffEqBase.solve(exprob::ExpectationProblem{F},
+                          expalg::MonteCarlo) where {F <: SystemMap}
+    d = distribution(exprob)
+    cov = input_cov(exprob)
+    S = mapping(exprob)
+    g = observable(exprob)
+
+    prob_func = function (prob, i, repeat)
+        u0, p = cov(rand(d), prob.u0, prob.p)
+        remake(prob, u0 = u0, p = p)
+    end
+
+    output_func(sol, i) = (g(sol, sol.prob.p), false)
+
+    monte_prob = EnsembleProblem(S.prob;
+                                 output_func = output_func,
+                                 prob_func = prob_func)
+    sol = solve(monte_prob, S.args...; trajectories = expalg.trajectories, S.kwargs...)
+    ExpectationSolution(mean(sol.u),nothing,nothing)
+end
+
+# Solve Koopman expectation
+function DiffEqBase.solve(prob::ExpectationProblem, expalg::Koopman, args...;
+                          maxiters = 1000000,
+                          batch = 0,
+                          quadalg = HCubatureJL(),
+                          ireltol = 1e-2, iabstol = 1e-2,
+                          kwargs...) where {A <: AbstractExpectationADAlgorithm}
+    integrand = build_integrand(prob, expalg, Val(batch > 1))
+    lb, ub = extrema(prob.d)
+
+    sol = integrate(quadalg, expalg.sensealg, integrand, lb, ub, prob.params;
+                    reltol = ireltol, abstol = iabstol, maxiters = maxiters,
+                    nout = prob.nout, batch = batch,
+                    kwargs...)
+
+    return ExpectationSolution(sol.u,sol.resid,sol)
+end
+
+# Integrate function to test new Adjoints, will need to roll up to Integrals.jl
+function integrate(quadalg, adalg::AbstractExpectationADAlgorithm, f, lb::TB, ub::TB, p;
+                   nout = 1, batch = 0,
+                   kwargs...) where {TB}
+    #TODO check batch iip type stability w/ IntegralProblem{XXXX}
+    prob = IntegralProblem{batch > 1}(f, lb, ub, p; nout = nout, batch = batch)
+    solve(prob, quadalg; kwargs...)
+end
+
+# defines adjoint via ∫∂/∂p f(x,p) dx
+Zygote.@adjoint function integrate(quadalg, adalg::NonfusedAD, f::F, lb::T, ub::T,
+                                   params::P;
+                                   nout = 1, batch = 0, norm = norm,
+                                   kwargs...) where {F, T, P}
+    primal = integrate(quadalg, adalg, f, lb, ub, params;
+                       norm = norm, nout = nout, batch = batch,
+                       kwargs...)
+
+    function integrate_pullbacks(Δ)
+        function dfdp(x, params)
+            _, back = Zygote.pullback(p -> f(x, p), params)
+            back(Δ)[1]
+        end
+        ∂p = integrate(quadalg, adalg, dfdp, lb, ub, params;
+                       norm = norm, nout = nout * length(params), batch = batch,
+                       kwargs...)
+        # ∂lb = -f(lb,params)  #needs correct for dim > 1
+        # ∂ub = f(ub,params)
+        return nothing, nothing, nothing, nothing, nothing, ∂p
+    end
+    primal, integrate_pullbacks
+end
+
+# defines adjoint via ∫[f(x,p; ∂/∂p f(x,p)] dx, ie it fuses the primal, post the primal calculation
+# has flag to only compute quad norm with respect to only the primal in the pull-back. Gives same quadrature points as doing forwarddiff
+Zygote.@adjoint function integrate(quadalg, adalg::PostfusedAD, f::F, lb::T, ub::T,
+                                   params::P;
+                                   nout = 1, batch = 0, norm = norm,
+                                   kwargs...) where {F, T, P}
+    primal = integrate(quadalg, adalg, f, lb, ub, params;
+                       norm = norm, nout = nout, batch = batch,
+                       kwargs...)
+
+    _norm = adalg.norm_partials ? norm : primalnorm(nout, norm)
+
+    function integrate_pullbacks(Δ)
+        function dfdp(x, params)
+            y, back = Zygote.pullback(p -> f(x, p), params)
+            [y; back(Δ)[1]]   #TODO need to match proper arrray type? promote_type???
+        end
+        ∂p = integrate(quadalg, adalg, dfdp, lb, ub, params;
+                       norm = _norm, nout = nout + nout * length(params), batch = batch,
+                       kwargs...)
+        return nothing, nothing, nothing, nothing, nothing, @view ∂p[(nout + 1):end]
+    end
+    primal, integrate_pullbacks
+end
+
+# Fuses primal and partials prior to pullback, I doubt this will stick around based on required system evals.
+Zygote.@adjoint function integrate(quadalg, adalg::PrefusedAD, f::F, lb::T, ub::T,
+                                   params::P;
+                                   nout = 1, batch = 0, norm = norm,
+                                   kwargs...) where {F, T, P}
+    # from Seth Axen via Slack
+    # Does not work w/ ArrayPartition unless with following hack
+    # Base.similar(A::ArrayPartition, ::Type{T}, dims::NTuple{N,Int}) where {T,N} = similar(Array(A), T, dims)
+    # TODO add ArrayPartition similar fix upstream, see https://github.com/SciML/RecursiveArrayTools.jl/issues/135
+    ∂f_∂params(x, params) = only(Zygote.jacobian(p -> f(x, p), params))
+    f_augmented(x, params) = [f(x, params); ∂f_∂params(x, params)...] #TODO need to match proper arrray type? promote_type???
+    _norm = adalg.norm_partials ? norm : primalnorm(nout, norm)
+
+    res = integrate(quadalg, adalg, f_augmented, lb, ub, params;
+                    norm = _norm, nout = nout + nout * length(params), batch = batch,
+                    kwargs...)
+    primal = first(res)
+    function integrate_pullback(Δy)
+        ∂params = Δy .* conj.(@view(res[(nout + 1):end]))
+        return nothing, nothing, nothing, nothing, nothing, ∂params
+    end
+    primal, integrate_pullback
+end
+
+# define norm function based only on primal part of fused integrand
+function primalnorm(nout, fnorm)
+    x -> fnorm(@view x[1:nout])
+end

src/expectation/problem_types.jl

@@ -0,0 +1,39 @@
+abstract type AbstractUncertaintyProblem end
+
+struct ExpectationProblem{TS, TG, TH, TF, TP} <: AbstractUncertaintyProblem
+    # defines ∫ g(S(h(x,u0,p)))*f(x)dx
+    # 𝕏 = uncertainty space, 𝕌 = Initial condition space, ℙ = model parameter space,
+    S::TS  # mapping,                 S: 𝕌 × ℙ → 𝕌
+    g::TG  # observable(output_func), g: 𝕌 × ℙ → ℝⁿᵒᵘᵗ
+    h::TH  # cov(input_func),         h: 𝕏 × 𝕌 × ℙ → 𝕌 × ℙ
+    d::TF  # distribution,            pdf(d,x): 𝕏 → ℝ
+    params::TP
+    nout::Int
+end
+
+# Constructor for general maps/functions
+function ExpectationProblem(g, pdist, params; nout = 1)
+    h(x, u, p) = x, p
+    S(x, p) = x
+    ExpectationProblem(S, g, h, pdist, params, nout)
+end
+
+# Constructor for DEProblems
+function ExpectationProblem(sm::SystemMap, g, h, d; nout = 1)
+    ExpectationProblem(sm, g, h, d,
+        ArrayPartition(deepcopy(sm.prob.u0),deepcopy(sm.prob.p)),
+        nout)
+end
+
+distribution(prob::ExpectationProblem) = prob.d
+mapping(prob::ExpectationProblem) = prob.S
+observable(prob::ExpectationProblem) = prob.g
+input_cov(prob::ExpectationProblem) = prob.h
+parameters(prob::ExpectationProblem) = prob.params
+
+##
+# struct CentralMomentProblem
+#     ns::NTuple{Int,N}
+#     altype::Union{NestedExpectation, BinomialExpansion} #Should rely be in solve
+#     exp_prob::ExpectationProblem
+# end

src/expectation/solution_types.jl

@@ -0,0 +1,5 @@
+struct ExpectationSolution{uType,R,O}
+    u::uType
+    resid::R
+    original::O
+end

src/expectation/system_utils.jl

@@ -0,0 +1,14 @@
+#Callable wrapper for DE solves. Enables seperation of args/kwargs...
+struct SystemMap{DT<:DiffEqBase.DEProblem,A,K}
+    prob::DT
+    args::A
+    kwargs::K
+end
+SystemMap(prob, args...; kwargs...) = SystemMap(prob, args, kwargs)
+
+function (sm::SystemMap{DT})(u0,p) where DT
+    prob::DT = remake(sm.prob, 
+                    u0 = convert(typeof(sm.prob.u0),u0),
+                    p = convert(typeof(sm.prob.p), p))
+    solve(prob, sm.args...; sm.kwargs...)
+end