[WIP] Bayesian PINN solver for ODE

Project.toml

@@ -5,6 +5,7 @@ version = "5.7.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
@@ -18,6 +19,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
 IntegralsCubature = "c31f79ba-6e32-46d4-a52f-182a8ac42a54"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
@@ -33,6 +35,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
+Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 

src/NeuralPDE.jl

@@ -48,6 +48,10 @@ include("transform_inf_integral.jl")
 include("discretize.jl")
 include("neural_adapter.jl")
 
+# fixes #682
+include("turing_MCMC.jl")
+include("advancedHMC_MCMC.jl")
+
 export NNODE, TerminalPDEProblem, NNPDEHan, NNPDENS, NNRODE,
        KolmogorovPDEProblem, NNKolmogorov, NNStopping, ParamKolmogorovPDEProblem,
        KolmogorovParamDomain, NNParamKolmogorov,
@@ -62,4 +66,7 @@ export NNODE, TerminalPDEProblem, NNPDEHan, NNPDENS, NNRODE,
        MiniMaxAdaptiveLoss,
        LogOptions
 
+#fixes #682
+export bayesian_pinn_ode, ahmc_bayesian_pinn_ode
+
 end # module

src/advancedHMC_MCMC.jl

@@ -0,0 +1,215 @@
+using AdvancedHMC, ForwardDiff, LogDensityProblems
+using LinearAlgebra
+
+mutable struct LogTargetDensity{C, S}
+    dim::Int
+    prob::DiffEqBase.DEProblem
+    chain::C
+    st::S
+    dataset::Tuple{AbstractVector, AbstractVector}
+    priorsNN::Tuple{Float64, Float64}
+    phystd::Float64
+    l2std::Float64
+    autodiff::Bool
+    physdt::Float64
+
+    function LogTargetDensity(dim, prob, chain::Optimisers.Restructure, st, dataset,
+                              priorsNN, phystd, l2std, autodiff, physdt)
+        new{typeof(chain), Nothing}(dim, prob, chain, nothing,
+                                    dataset, priorsNN,
+                                    phystd, l2std, autodiff,
+                                    physdt)
+    end
+    function LogTargetDensity(dim, prob, chain::Lux.AbstractExplicitLayer, st, dataset,
+                              priorsNN, phystd, l2std, autodiff, physdt)
+        new{typeof(chain), typeof(st)}(dim, prob, re, st,
+                                       dataset, priorsNN,
+                                       phystd, l2std, autodiff,
+                                       physdt)
+    end
+end
+
+function LogDensityProblems.logdensity(Tar::LogTargetDensity, θ)
+    return physloglikelihood(Tar, θ) + L2LossData(Tar, θ) + priorweights(Tar, θ)
+end
+
+LogDensityProblems.dimension(Tar::LogTargetDensity) = Tar.dim
+
+function LogDensityProblems.capabilities(::LogTargetDensity)
+    LogDensityProblems.LogDensityOrder{0}()
+end
+
+function generate_Tar(chain::Lux.AbstractExplicitLayer, init_params)
+    θ, st = Lux.setup(Random.default_rng(), chain)
+    return ComponentArrays.ComponentArray(init_params), chain, st
+end
+
+function generate_Tar(chain::Lux.AbstractExplicitLayer, init_params::Nothing)
+    θ, st = Lux.setup(Random.default_rng(), chain)
+    return ComponentArrays.ComponentArray(θ), chain, st
+end
+
+function generate_Tar(chain::Flux.Chain, init_params)
+    θ, re = Flux.destructure(chain)
+    # if init_params==nothing
+    #     θ = collect(Float64, θ)
+    #     return θ, re, nothing
+    # else 
+    #     return init_params, re, nothing
+    # end
+    return init_params, re, nothing
+end
+function generate_Tar(chain::Flux.Chain, init_params::Nothing)
+    θ, re = Flux.destructure(chain)
+    # find_good_stepsize takes only float64?
+    θ = collect(Float64, θ)
+    return θ, re, nothing
+end
+
+# nn OUTPUT AT t
+function (f::LogTargetDensity{C, S})(t::AbstractVector,
+                                     θ) where {C <: Optimisers.Restructure, S}
+    # f.prob.u0 .+ (t .- f.prob.tspan[1]) .* vec(f.chain(θ)(t'))
+    f.prob.u0 .+ (t' .- f.prob.tspan[1]) .* f.chain(θ)(adapt(parameterless_type(θ), t'))
+    # f.prob.u0 .+ (t' .- f.prob.tspan[1]) .* f.chain(θ)(t')
+end
+
+function (f::LogTargetDensity{C, S})(t::AbstractVector,
+                                     θ) where {C <: Lux.AbstractExplicitLayer, S}
+    # Batch via data as row vectors
+    y, st = f.chain(adapt(parameterless_type(ComponentArrays.getdata(θ)), t'), θ, f.st)
+    ChainRulesCore.@ignore_derivatives f.st = st
+    f.u0 .+ (t' .- f.t0) .* y
+end
+
+function (f::LogTargetDensity{C, S})(t::Number,
+                                     θ) where {C <: Optimisers.Restructure, S}
+    f.prob.u0 + (t - f.prob.tspan[1]) * first(f.chain(θ)(adapt(parameterless_type(θ), [t])))
+end
+
+using ForwardDiff: Chunk
+# ODE DU/DX
+function NNodederi(phi::LogTargetDensity, t::AbstractVector, θ, autodiff::Bool)
+    if autodiff
+        # diag(ForwardDiff.jacobian(t -> phi(t, θ), t))
+        [ForwardDiff.derivative(ti -> phi(ti, θ), ti) for ti in t]
+    else
+        (phi(t .+ sqrt(eps(eltype(t))), θ) - phi(t, θ)) ./ sqrt(eps(eltype(t)))
+    end
+end
+
+function physloglikelihood(Tar::LogTargetDensity, θ)
+    p = Tar.prob.p
+    f = Tar.prob.f
+    var = Tar.phystd^2
+    autodiff = Tar.autodiff
+    dt = Tar.physdt
+    t = collect(Float64, Tar.prob.tspan[1]:dt:Tar.prob.tspan[2])
+
+    # compare derivatives
+    out = Tar(t, θ)
+    # print(size(out))
+    physsol = [f(out[i], p, t[i]) for i in eachindex(out)]
+    nnsol = NNodederi(Tar, t, θ, autodiff)
+
+    # distribution's mean is forwarddiff diag(jacobian)
+    # if autodiff
+    #     nnsol = vec(diag(nnsol))
+    # end
+    # typeof(nnsol), size(nnsol), typeof(physsol), size(physsol)
+    n = length(nnsol)
+    return logpdf(MvNormal(nnsol, Diagonal(var .* ones(n))), physsol)
+end
+
+# standard MvNormal Dist Assume
+function L2LossData(Tar::LogTargetDensity, θ)
+    nn = vec(Tar(Tar.dataset[2], θ))
+    n = length(nn)
+    var = Tar.l2std^2
+    return logpdf(MvNormal(nn, Diagonal(var .* ones(n))), Tar.dataset[1])
+end
+
+function priorweights(Tar::LogTargetDensity, θ)
+    params = Tar.priorsNN
+    return logpdf(MvNormal(θ, Diagonal(params[2]^2 .* ones(length(θ)))),
+                  params[1] * ones(length(θ)))
+end
+
+# dataset would be (x̂,t)
+# priors: pdf for W,b + pdf for ODE params
+function ahmc_bayesian_pinn_ode(prob::DiffEqBase.DEProblem, chain::Flux.Chain,
+                                dataset::Tuple{AbstractVector, AbstractVector};
+                                init_params = nothing,
+                                draw_samples = 1000, l2std = 0.08,
+                                phystd = 0.08, priorsNN = (0, 2), autodiff = false,
+                                physdt = 1 / 20.0f0,
+                                Proposal = AdvancedHMC.NUTS{MultinomialTS,
+                                                            GeneralisedNoUTurn},
+                                Adaptor = StanHMCAdaptor, targetacceptancerate = 0.8,
+                                Integrator = Leapfrog, Metric = DiagEuclideanMetric)
+    # NN parameter prior mean and variance(PriorsNN must be a tuple)
+    if isinplace(prob)
+        throw(error("The BPINN ODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t)."))
+    end
+
+    if chain isa Lux.AbstractExplicitLayer || chain isa Flux.Chain
+        initial_θ, recon, st = generate_Tar(chain, init_params)
+    else
+        error("Only Lux.AbstractExplicitLayer and Flux.Chain neural networks are supported")
+    end
+
+    # adding ode parameter estimation?
+    nparameters = length(initial_θ)
+    ℓπ = LogTargetDensity(nparameters, prob, recon, st, dataset, priorsNN,
+                          phystd, l2std, autodiff, physdt)
+
+    #  [add f(t,θ) for t being a number]
+    # try
+    #     ℓπ(t0, initial_θ)
+    # catch err
+    #     if isa(err, DimensionMismatch)
+    #         throw(DimensionMismatch("Dimensions of the initial u0 and chain should match"))
+    #     else
+    #         throw(err)
+    #     end
+    # end
+
+    # return physloglikelihood(ℓπ, initial_θ)
+    n_samples = draw_samples
+    metric = Metric(nparameters)
+    hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)
+    initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
+
+    # choices for integrators?
+    # [define n for JL and a for TL]
+    # if Integrator == JitteredLeapfrog(n)
+    #     integrator = JitteredLeapfrog(initial_ϵ, n)
+    # elseif Integrator == TemperedLeapfrog(a)
+    #     integrator == TemperedLeapfrog(initial_ϵ, a)
+    # else
+    #     integrator = Leapfrog(initial_ϵ)
+    # end
+
+    integrator = Leapfrog(initial_ϵ)
+    proposal = Proposal(integrator)
+    adaptor = Adaptor(MassMatrixAdaptor(metric),
+                      StepSizeAdaptor(targetacceptancerate, integrator))
+
+    samples, stats = sample(hamiltonian, proposal, initial_θ,
+                            n_samples, adaptor;
+                            progress = true)
+
+    # return a chain(basic chain),samples and stats
+    matrix_samples = hcat(samples...)
+    mcmc_chain = Chains(matrix_samples')
+    return mcmc_chain, samples, stats
+end
+
+# non vectorised functions(i noticed sampling time increase)
+# function NNodederi(phi::odeByNN, t::Number, θ, autodiff::Bool)
+#     if autodiff
+#         ForwardDiff.jacobian(t -> phi(t, θ), t)
+#     else
+#         (phi(t + sqrt(eps(typeof(t))), θ) - phi(t, θ)) / sqrt(eps(typeof(t)))
+#     end
+# end

src/turing_MCMC.jl

@@ -0,0 +1,125 @@
+using Turing, Distributions
+struct odeByNN{C, T, U}
+    chain::C
+    u0::U
+    t0::T
+
+    function odeByNN(re::Optimisers.Restructure, t, u0)
+        new{typeof(re), typeof(t), typeof(u0)}(re, t, u0)
+    end
+end
+
+function generate_phi(chain::Flux.Chain, t, u0, init_params::Nothing)
+    θ, re = Flux.destructure(chain)
+    odeByNN(re, t, u0), θ
+end
+
+# nn OUTPUT AT t
+function (f::odeByNN{C, T, U})(t::Number,
+    θ) where {C <: Optimisers.Restructure, T, U}
+    f.u0 + (t - f.t0) * first(f.chain(θ)(adapt(parameterless_type(θ), [t])))
+end
+
+function (f::odeByNN{C, T, U})(t::AbstractVector,
+    θ) where {C <: Optimisers.Restructure, T, U}
+    f.u0 .+ (t .- f.t0) .* f.chain(θ)(adapt(parameterless_type(θ), t'))
+end
+
+# ODE DU/DX
+function NNodederi(phi::odeByNN, t::Number, θ, autodiff::Bool)
+    if autodiff
+        ForwardDiff.jacobian(t -> phi(t, θ), t)
+    else
+        (phi(t + sqrt(eps(typeof(t))), θ) - phi(t, θ)) / sqrt(eps(typeof(t)))
+    end
+end
+
+function NNodederi(phi::odeByNN, t::AbstractVector, θ, autodiff::Bool)
+    if autodiff
+        ForwardDiff.jacobian(t -> phi(t, θ), t)
+    else
+        (phi(t .+ sqrt(eps(eltype(t))), θ) - phi(t, θ)) ./ sqrt(eps(eltype(t)))
+    end
+end
+
+function physloglikelihood(chain::Any, prob::DiffEqBase.DEProblem,
+    t::AbstractVector; var = 0.5)
+    u0 = prob.u0
+    t0 = t[1]
+    p = prob.p
+    f = prob.f
+    # let this be(will fix)
+    autodiff = false
+
+    phi, initparams = generate_phi(chain, t0, u0, nothing)
+
+    μ = vec([f(phi(t[i], initparams), p, u0) for i in eachindex(t)])
+    physsol = vec([NNodederi(phi, t[i], initparams, autodiff) for i in eachindex(t)])
+    # print(typeof(μ))
+    # print(typeof(physsol))
+    # To reduce heap allocations but some erros came up
+    # μ = similar(t)
+    # physsol = similar(t)
+    # μ = f(phi(t, initparams), p, u0)
+    # physsol = NNodederi(phi, t, initparams, autodiff)
+    # for i in eachindex(t)
+    #     μ[i] = f(phi(t[i], initparams), p, u0)
+    #     physsol[i] = NNodederi(phi, t[i], initparams, autodiff)
+    # end
+    return sum(abs2, (μ .- physsol) ./ (-2 * (var^2)))
+    # return loglikelihood(MvNormal(physsol - μ,
+    #   Diagonal(var .* ones(Float64, length(μ)))))
+end
+
+# dataset would be (x̂,t)
+# priors: pdf for W,b + pdf for ODE params
+function bayesian_pinn_ode(prob::DiffEqBase.DEProblem, chain, dataset;
+    sampling_strategy = Turing.NUTS(0.65), num_samples = 1000)
+    param_initial, recon = Flux.destructure(chain)
+    nparameters = length(param_initial)
+
+    if isinplace(prob)
+        throw(error("The BPINN ODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t)."))
+    end
+
+    alpha = 0.09
+    sig = 0.2
+
+    p = prob.p
+    f = prob.f
+    u0 = prob.u0
+    t0 = prob.tspan[1]
+    t = dataset[2]
+    # compare derivatives
+    phi, initparams = generate_phi(chain, t0, u0, nothing)
+    physsol = vec([f(phi(t[i], param_initial), p, t[i]) for i in eachindex(t)])
+
+    DynamicPPL.@model function bayes_pinn(dataset, physsol, phi)
+        # parameter estimation?
+
+        # prior for NN parameters(not included bias yet?) - P(Θ)
+        nnparameters ~ MvNormal(zeros(nparameters), Diagonal(sig .* ones(nparameters)))
+        preds = [phi(ti, nnparameters) for ti in dataset[2]]
+
+        # # likelihood for NN pred vs Equation satif - P(phys | Θ)
+        # if DynamicPPL.leafcontext(__context__) !== Turing.PriorContext()
+        #     Turing.@addlogprob! physloglikelihood(nn, prob, dataset[2], var = sig)
+        # end
+        physsol ~ MvNormal(vec(preds), Diagonal(sig .* ones(length(dataset[2]))))
+        # # likelihood for dataset vs NN pred  - P( X̄ | Θ)
+        dataset[1] ~ MvNormal(vec(preds), Diagonal(sig .* ones(length(dataset[2]))))
+    end
+
+    model = bayes_pinn(dataset, physsol, phi)
+    ch = sample(model, sampling_strategy, num_samples)
+    return ch
+end
+
+# ----------need speed up
+# the phase point struct
+# create custom distri?
+# using chain with updated parameters in physloglikelihood and L2LossData
+
+# ----------more code and compatibility
+# allow options for prior,likelihood distributions
+# add support for Lux chains