Added a separate package BayesianNeuralPDE.jl

Project.toml

@@ -22,6 +22,7 @@ IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"

lib/BayesianNeuralPDE/Project.toml

@@ -0,0 +1,30 @@
+name = "BayesianNeuralPDE"
+uuid = "3cea9122-e921-42ea-a9d7-c72fcb58ce53"
+authors = ["paramthakkar123 <[email protected]>"]
+version = "0.1.0"
+
+[deps]
+AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
+OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
+NeuralPDE = "315f7962-48a3-4962-8226-d0f33b1235f0"
+
+[compat]
+ChainRulesCore = "1.25.1"
+ConcreteStructs = "0.2.3"
+MonteCarloMeasurements = "1.4.3"
+Printf = "1.11.0"
+SciMLBase = "2.72.1"

lib/BayesianNeuralPDE/src/BPINN_ode.jl

@@ -0,0 +1,231 @@
+# HIGH level API for BPINN ODE solver
+
+"""
+	BNNODE(chain, kernel = HMC; strategy = nothing, draw_samples = 2000,
+		   priorsNNw = (0.0, 2.0), param = [nothing], l2std = [0.05],
+		   phystd = [0.05], phynewstd = [0.05], dataset = [nothing], physdt = 1 / 20.0,
+		   MCMCargs = (; n_leapfrog=30), nchains = 1, init_params = nothing,
+		   Adaptorkwargs = (; Adaptor = StanHMCAdaptor, targetacceptancerate = 0.8,
+							  Metric = DiagEuclideanMetric),
+		   Integratorkwargs = (Integrator = Leapfrog,), autodiff = false,
+		   progress = false, verbose = false)
+
+Algorithm for solving ordinary differential equations using a Bayesian neural network. This
+is a specialization of the physics-informed neural network which is used as a solver for a
+standard `ODEProblem`.
+
+!!! warn
+
+	Note that BNNODE only supports ODEs which are written in the out-of-place form, i.e.
+	`du = f(u,p,t)`, and not `f(du,u,p,t)`. If not declared out-of-place, then the BNNODE
+	will exit with an error.
+
+## Positional Arguments
+
+* `chain`: A neural network architecture, defined as a `Lux.AbstractLuxLayer`.
+* `kernel`: Choice of MCMC Sampling Algorithm. Defaults to `AdvancedHMC.HMC`
+
+## Keyword Arguments
+
+(refer `NeuralPDE.ahmc_bayesian_pinn_ode` keyword arguments.)
+
+## Example
+
+```julia
+linear = (u, p, t) -> -u / p[1] + exp(t / p[2]) * cos(t)
+tspan = (0.0, 10.0)
+u0 = 0.0
+p = [5.0, -5.0]
+prob = ODEProblem(linear, u0, tspan, p)
+linear_analytic = (u0, p, t) -> exp(-t / 5) * (u0 + sin(t))
+
+sol = solve(prob, Tsit5(); saveat = 0.05)
+u = sol.u[1:100]
+time = sol.t[1:100]
+x̂ = u .+ (u .* 0.2) .* randn(size(u))
+dataset = [x̂, time]
+
+chainlux = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1))
+
+alg = BNNODE(chainlux; draw_samples = 2000, l2std = [0.05], phystd = [0.05],
+			 priorsNNw = (0.0, 3.0), progress = true)
+
+sol_lux = solve(prob, alg)
+
+# with parameter estimation
+alg = BNNODE(chainlux; dataset, draw_samples = 2000, l2std = [0.05], phystd = [0.05],
+			 priorsNNw = (0.0, 10.0), param = [Normal(6.5, 0.5), Normal(-3, 0.5)],
+			 progress = true)
+
+sol_lux_pestim = solve(prob, alg)
+```
+
+## Solution Notes
+
+Note that the solution is evaluated at fixed time points according to the strategy chosen.
+ensemble solution is evaluated and given at steps of `saveat`.
+Dataset should only be provided when ODE parameter Estimation is being done.
+The neural network is a fully continuous solution so `BPINNsolution`
+is an accurate interpolation (up to the neural network training result). In addition, the
+`BPINNstats` is returned as `sol.fullsolution` for further analysis.
+
+## References
+
+Liu Yanga, Xuhui Menga, George Em Karniadakis. "B-PINNs: Bayesian Physics-Informed Neural
+Networks for Forward and Inverse PDE Problems with Noisy Data".
+
+Kevin Linka, Amelie Schäfer, Xuhui Meng, Zongren Zou, George Em Karniadakis, Ellen Kuhl
+"Bayesian Physics Informed Neural Networks for real-world nonlinear dynamical systems".
+"""
+@concrete struct BNNODE <: NeuralPDEAlgorithm
+	chain <: AbstractLuxLayer
+	kernel::Any
+	strategy <: Union{Nothing, AbstractTrainingStrategy}
+	draw_samples::Int
+	priorsNNw::Tuple{Float64, Float64}
+	param <: Union{Nothing, Vector{<:Distribution}}
+	l2std::Vector{Float64}
+	phystd::Vector{Float64}
+	phynewstd::Vector{Float64}
+	dataset <: Union{Vector{Nothing}, Vector{<:Vector{<:AbstractFloat}}}
+	physdt::Float64
+	MCMCkwargs <: NamedTuple
+	nchains::Int
+	init_params <: Union{Nothing, <:NamedTuple, Vector{<:AbstractFloat}}
+	Adaptorkwargs <: NamedTuple
+	Integratorkwargs <: NamedTuple
+	numensemble::Int
+	estim_collocate::Bool
+	autodiff::Bool
+	progress::Bool
+	verbose::Bool
+end
+
+function BNNODE(chain, kernel = HMC; strategy = nothing, draw_samples = 1000,
+	priorsNNw = (0.0, 2.0), param = nothing, l2std = [0.05], phystd = [0.05],
+	phynewstd = [0.05], dataset = [nothing], physdt = 1 / 20.0,
+	MCMCkwargs = (n_leapfrog = 30,), nchains = 1, init_params = nothing,
+	Adaptorkwargs = (Adaptor = StanHMCAdaptor,
+		Metric = DiagEuclideanMetric, targetacceptancerate = 0.8),
+	Integratorkwargs = (Integrator = Leapfrog,),
+	numensemble = floor(Int, draw_samples / 3),
+	estim_collocate = false, autodiff = false, progress = false, verbose = false)
+	chain isa AbstractLuxLayer || (chain = FromFluxAdaptor()(chain))
+	return BNNODE(chain, kernel, strategy, draw_samples, priorsNNw, param, l2std, phystd,
+		phynewstd, dataset, physdt, MCMCkwargs, nchains, init_params, Adaptorkwargs,
+		Integratorkwargs, numensemble, estim_collocate, autodiff, progress, verbose)
+end
+
+"""
+Contains `ahmc_bayesian_pinn_ode()` function output:
+
+1. A MCMCChains.jl chain object for sampled parameters.
+2. The set of all sampled parameters.
+3. Statistics like:
+	- n_steps
+	- acceptance_rate
+	- log_density
+	- hamiltonian_energy
+	- hamiltonian_energy_error
+	- numerical_error
+	- step_size
+	- nom_step_size
+"""
+@concrete struct BPINNstats
+	mcmc_chain::Any
+	samples::Any
+	statistics::Any
+end
+
+"""
+BPINN Solution contains the original solution from AdvancedHMC.jl sampling (BPINNstats
+contains fields related to that).
+
+1. `ensemblesol` is the Probabilistic Estimate (MonteCarloMeasurements.jl Particles type) of
+   Ensemble solution from All Neural Network's (made using all sampled parameters) output's.
+2. `estimated_nn_params` - Probabilistic Estimate of NN params from sampled weights, biases.
+3. `estimated_de_params` - Probabilistic Estimate of DE params from sampled unknown DE
+   parameters.
+"""
+@concrete struct BPINNsolution
+	original <: BPINNstats
+	ensemblesol::Any
+	estimated_nn_params::Any
+	estimated_de_params::Any
+	timepoints::Any
+end
+
+function SciMLBase.__solve(prob::SciMLBase.ODEProblem, alg::BNNODE, args...; dt = nothing,
+	timeseries_errors = true, save_everystep = true, adaptive = false,
+	abstol = 1.0f-6, reltol = 1.0f-3, verbose = false, saveat = 1 / 50.0,
+	maxiters = nothing)
+	(; chain, param, strategy, draw_samples, numensemble, verbose) = alg
+
+	# ahmc_bayesian_pinn_ode needs param=[] for easier vcat operation for full vector of parameters
+	param = param === nothing ? [] : param
+	strategy = strategy === nothing ? GridTraining : strategy
+
+	@assert alg.draw_samples ≥ 0 "Number of samples to be drawn has to be >=0."
+
+	mcmcchain, samples, statistics = ahmc_bayesian_pinn_ode(
+		prob, chain; strategy, alg.dataset, alg.draw_samples, alg.init_params,
+		alg.physdt, alg.l2std, alg.phystd, alg.phynewstd,
+		alg.priorsNNw, param, alg.nchains, alg.autodiff,
+		Kernel = alg.kernel, alg.Adaptorkwargs, alg.Integratorkwargs,
+		alg.MCMCkwargs, alg.progress, alg.verbose, alg.estim_collocate)
+
+	fullsolution = BPINNstats(mcmcchain, samples, statistics)
+	ninv = length(param)
+	t = collect(eltype(saveat), prob.tspan[1]:saveat:prob.tspan[2])
+
+	θinit, st = LuxCore.setup(Random.default_rng(), chain)
+	θ = [vector_to_parameters(samples[i][1:(end-ninv)], θinit)
+		 for i in (draw_samples-numensemble):draw_samples]
+
+	luxar = [chain(t', θ[i], st)[1] for i in 1:numensemble]
+	# only need for size
+	θinit = collect(ComponentArray(θinit))
+
+	# constructing ensemble predictions
+	ensemblecurves = Vector{}[]
+	# check if NN output is more than 1
+	numoutput = size(luxar[1])[1]
+	if numoutput > 1
+		# Initialize a vector to store the separated outputs for each output dimension
+		output_matrices = [Vector{Vector{Float32}}() for _ in 1:numoutput]
+
+		# Loop through each element in `luxar`
+		for element in luxar
+			for i in 1:numoutput
+				push!(output_matrices[i], element[i, :])  # Append the i-th output (i-th row) to the i-th output_matrices
+			end
+		end
+
+		for r in 1:numoutput
+			ensem_r = hcat(output_matrices[r]...)'
+			ensemblecurve_r = prob.u0[r] .+
+							  [Particles(ensem_r[:, i]) for i in 1:length(t)] .*
+							  (t .- prob.tspan[1])
+			push!(ensemblecurves, ensemblecurve_r)
+		end
+
+	else
+		ensemblecurve = prob.u0 .+
+						[Particles(reduce(vcat, luxar)[:, i]) for i in 1:length(t)] .*
+						(t .- prob.tspan[1])
+		push!(ensemblecurves, ensemblecurve)
+	end
+
+	nnparams = length(θinit)
+	estimnnparams = [Particles(reduce(hcat, samples[(end-numensemble):end])[i, :])
+					 for i in 1:nnparams]
+
+	if ninv == 0
+		estimated_params = [nothing]
+	else
+		estimated_params = [Particles(reduce(hcat, samples[(end-numensemble):end])[i, :])
+							for i in (nnparams+1):(nnparams+ninv)]
+	end
+
+	return BPINNsolution(fullsolution, ensemblecurves, estimnnparams, estimated_params, t)
+end

lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl

@@ -0,0 +1,26 @@
+module BayesianNeuralPDE
+
+using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux,
+	AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements
+using Printf: @printf
+using ConcreteStructs: @concrete
+using NeuralPDE: PhysicsInformedNN
+using SciMLBase: SciMLBase
+using ChainRulesCore: ChainRulesCore, @non_differentiable, @ignore_derivatives
+using LogDensityProblems: LogDensityProblems
+
+abstract type AbstractPINN end
+
+abstract type AbstractTrainingStrategy end
+abstract type NeuralPDEAlgorithm <: SciMLBase.AbstractODEAlgorithm end
+
+include("advancedHMC_MCMC.jl")
+include("pinn_types.jl")
+include("BPINN_ode.jl")
+include("discretize.jl")
+include("PDE_BPINN.jl")
+
+export BNNODE, ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde
+export BPINNsolution, BayesianPINN
+
+end