BayesianPINN Discretizer for PDESystems

Using the Bayesian PINN solvers, we can solve general nonlinear PDEs, ODEs and also simultaneously perform parameter estimation on them.

Note: The BPINN PDE solver also works for ODEs defined using ModelingToolkit, ModelingToolkit.jl PDESystem documentation. Despite this, the ODE specific BPINN solver BNNODErefer exists and uses NeuralPDE.ahmc_bayesian_pinn_ode at a lower level.

BayesianPINN Discretizer for PDESystems and lower level Bayesian PINN Solver calls for PDEs and ODEs.

NeuralPDE.BayesianPINNType
BayesianPINN(args...; dataset = nothing, kwargs...)

A discretize algorithm for the ModelingToolkit PDESystem interface, which transforms a PDESystem into a likelihood function used for HMC based Posterior Sampling Algorithms AdvancedHMC.jl which is later optimized upon to give the Solution Distribution of the PDE, using the Physics-Informed Neural Networks (PINN) methodology.

All positional arguments and keyword arguments are passed to PhysicsInformedNN except the ones mentioned below.

Keyword Arguments

  • dataset: A vector of matrix, each matrix for ith dependant variable and first col in matrix is for dependant variables, remaining columns for independent variables. Needed for inverse problem solving.
source
NeuralPDE.ahmc_bayesian_pinn_odeFunction
ahmc_bayesian_pinn_ode(prob, chain; strategy = GridTraining, dataset = [],
                       init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0f0,
                       l2std = [0.05], phystd = [0.05], phynewstd = (ode_params)->[0.05], priorsNNw = (0.0, 2.0),
                       param = [], nchains = 1, autodiff = false, Kernel = HMC,
                       Adaptorkwargs = (Adaptor = StanHMCAdaptor,
                           Metric = DiagEuclideanMetric, targetacceptancerate = 0.8),
                       Integratorkwargs = (Integrator = Leapfrog,),
                       MCMCkwargs = (n_leapfrog = 30,), progress = false,
                       verbose = false)
Warn

Note that ahmc_bayesian_pinn_ode() 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 ahmc_bayesian_pinn_ode() will exit with an error.

Example

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)

### CREATE DATASET (Necessity for accurate Parameter estimation)
sol = solve(prob, Tsit5(); saveat = 0.05)
u = sol.u[1:100]
time = sol.t[1:100]

### dataset and BPINN create
x̂ = collect(Float64, Array(u) + 0.05 * randn(size(u)))
dataset = [x̂, time, 0.05 .* ones(length(time))]

chain1 = Lux.Chain(Lux.Dense(1, 5, tanh), Lux.Dense(5, 5, tanh), Lux.Dense(5, 1)

### simply solving ode here hence better to not pass dataset(uses ode params specified in prob)
fh_mcmc_chain1, fhsamples1, fhstats1 = ahmc_bayesian_pinn_ode(prob, chain1,
                                                            dataset = dataset,
                                                            draw_samples = 1500,
                                                            l2std = [0.05],
                                                            phystd = [0.05],
                                                            priorsNNw = (0.0,3.0))

### solving ode + estimating parameters hence dataset needed to optimize parameters upon + Pior Distributions for ODE params
fh_mcmc_chain2, fhsamples2, fhstats2 = ahmc_bayesian_pinn_ode(prob, chain1,
                                                            dataset = dataset,
                                                            draw_samples = 1500,
                                                            l2std = [0.05],
                                                            phystd = [0.05],
                                                            priorsNNw = (0.0,3.0),
                                                            param = [Normal(6.5,0.5), Normal(-3,0.5)])

NOTES

Dataset is required for accurate Parameter estimation in Inverse Problems. Incase you are only solving Non parametric ODE Equations for a solution, do not provide a dataset.

Positional Arguments

  • prob: ODEProblem(out of place and the function signature should be f(u,p,t).
  • chain: Lux Neural Netork which would be made the Bayesian PINN.

Keyword Arguments

  • strategy: The training strategy used to choose the points for the evaluations. By default GridTraining is used with given physdt discretization.
  • dataset: Is either an empty Vector or a nested Vector of the form [x̂, t, W] where are dependant variable observations, t are time points and W are quadrature weights for domain. The dataset is used to compute the L2 loss against the data and also for the Data Duadrature loss function. For multiple dependant variables, there will be multiple vectors with the last two vectors in dataset still being for t, W. Is empty by default assuming a forward problem is being solved.
  • init_params: initial parameter values for BPINN (ideally for multiple chains different initializations preferred)
  • nchains: number of chains you want to sample
  • draw_samples: number of samples to be drawn in the MCMC algorithms (warmup samples are ~2/3 of draw samples)
  • l2std: standard deviation of BPINN prediction against L2 losses/Dataset
  • phystd: standard deviation of BPINN prediction against Chosen Underlying ODE System
  • phynewstd: A function that gives the standard deviation of the Data Quadrature loss function at each iteration. It takes the ODE parameters as input and returns a vector of standard deviations. Is (ode_params) -> [0.05] by default.
  • priorsNNw: Tuple of (mean, std) for BPINN Network parameters. Weights and Biases of BPINN are Normal Distributions by default.
  • param: Vector of chosen ODE parameters Distributions in case of Inverse problems.
  • autodiff: Boolean Value for choice of Derivative Backend(default is numerical)
  • physdt: Timestep for approximating ODE in it's Time domain. (1/20.0 by default)
  • Kernel: Choice of MCMC Sampling Algorithm (AdvancedHMC.jl implementations HMC/NUTS/HMCDA)
  • Integratorkwargs: Integrator, jitter_rate, tempering_rate. Refer: https://turinglang.org/AdvancedHMC.jl/stable/
  • Adaptorkwargs: Adaptor, Metric, targetacceptancerate. Refer: https://turinglang.org/AdvancedHMC.jl/stable/ Note: Target percentage (in decimal) of iterations in which the proposals are accepted (0.8 by default)
  • MCMCargs: A NamedTuple containing all the chosen MCMC kernel's (HMC/NUTS/HMCDA) Arguments, as follows :
    • n_leapfrog: number of leapfrog steps for HMC
    • δ: target acceptance probability for NUTS and HMCDA
    • λ: target trajectory length for HMCDA
    • max_depth: Maximum doubling tree depth (NUTS)
    • Δ_max: Maximum divergence during doubling tree (NUTS)
    Refer: https://turinglang.org/AdvancedHMC.jl/stable/
  • estim_collocate: A boolean value to indicate whether to use the Data Quadrature loss function or not. This is only relevant for ODE parameter estimation.
  • progress: controls whether to show the progress meter or not.
  • verbose: controls the verbosity. (Sample call args in AHMC)
Warning

AdvancedHMC.jl is still developing convenience structs so might need changes on new releases.

source
NeuralPDE.ahmc_bayesian_pinn_pdeFunction
ahmc_bayesian_pinn_pde(pde_system, discretization;
    draw_samples = 1000, bcstd = [0.01], l2std = [0.05], phystd = [0.05],
    phynewstd = [0.05], priorsNNw = (0.0, 2.0), param = [], nchains = 1,
    Kernel = HMC(0.1, 30), Adaptorkwargs = (Adaptor = StanHMCAdaptor,
        Metric = DiagEuclideanMetric, targetacceptancerate = 0.8),
    Integratorkwargs = (Integrator = Leapfrog,), saveats = [1 / 10.0],
    numensemble = floor(Int, draw_samples / 3), progress = false, verbose = false)

NOTES

  • Dataset is required for accurate Parameter estimation + solving equations.
  • Returned solution is a BPINNsolution consisting of Ensemble solution, estimated PDE and NN parameters for chosen saveats grid spacing and last n = numensemble samples in Chain. the complete set of samples in the MCMC chain is returned as fullsolution, refer BPINNsolution for more details.

Positional Arguments

  • pde_system: ModelingToolkit defined PDE equation or system of equations.
  • discretization: BayesianPINN discretization for the given pde_system, Neural Network and training strategy.

Keyword Arguments

  • draw_samples: number of samples to be drawn in the MCMC algorithms (warmup samples are ~2/3 of draw samples)
  • bcstd: Vector of standard deviations of BPINN prediction against Initial/Boundary Condition equations.
  • l2std: Vector of standard deviations of BPINN prediction against L2 losses/Dataset for each dependant variable of interest.
  • phystd: Vector of standard deviations of BPINN prediction against Chosen Underlying PDE equations.
  • phynewstd: Vector of standard deviations of the Data Quadrature loss term.
  • priorsNNw: Tuple of (mean, std) for BPINN Network parameters. Weights and Biases of BPINN are Normal Distributions by default.
  • param: Vector of chosen PDE's parameter's Distributions in case of Inverse problems.
  • nchains: number of chains you want to sample.
  • Kernel: Choice of MCMC Sampling Algorithm object HMC/NUTS/HMCDA (AdvancedHMC.jl implementations).
  • Adaptorkwargs: Adaptor, Metric, targetacceptancerate. Refer: https://turinglang.org/AdvancedHMC.jl/stable/. Note: Target percentage(in decimal) of iterations in which the proposals are accepted (0.8 by default).
  • Integratorkwargs: Integrator, jitter_rate, tempering_rate. Refer: https://turinglang.org/AdvancedHMC.jl/stable/
  • saveats: Grid spacing for each independent variable for evaluation of ensemble solution, estimated parameters.
  • numensemble: Number of last samples to take for creation of ensemble solution, estimated parameters.
  • progress: controls whether to show the progress meter or not.
  • verbose: controls the verbosity. (Sample call args in AHMC).
Warning

AdvancedHMC.jl is still developing convenience structs so might need changes on new releases.

source

symbolic_discretize for BayesianPINN and lower level interface.

SciMLBase.symbolic_discretizeMethod
prob = symbolic_discretize(pde_system::PDESystem, discretization::AbstractPINN)

symbolic_discretize is the lower level interface to discretize for inspecting internals. It transforms a symbolic description of a ModelingToolkit-defined PDESystem into a PINNRepresentation which holds the pieces required to build an OptimizationProblem for Optimization.jl or a Likelihood Function used for HMC based Posterior Sampling Algorithms AdvancedHMC.jl which is later optimized upon to give Solution or the Solution Distribution of the PDE.

For more information, see discretize and PINNRepresentation.

source
NeuralPDE.BPINNstatsType

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
    • hamiltonianenergyerror
    • numerical_error
    • step_size
    • nomstepsize
source
NeuralPDE.BPINNsolutionType

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.
source