From 41fd26647b1d9f6de0de3bc87a951bb4888d940d Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Tue, 28 Jan 2025 07:54:56 +0530 Subject: [PATCH 01/15] Added a separate package BayesianNeuralPDE.jl --- lib/BayesianNeuralPDE/Project.toml | 4 + .../BayesianNeuralPDE/src}/BPINN_ode.jl | 2 +- .../src/BayesianNeuralPDE.jl | 11 + .../BayesianNeuralPDE/src}/PDE_BPINN.jl | 2 +- lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl | 240 +++++++++++++++++ lib/BayesianNeuralPDE/src/pinn_types.jl | 33 +++ .../test}/BPINN_PDE_tests.jl | 0 .../BayesianNeuralPDE/test}/BPINN_tests.jl | 0 lib/BayesianNeuralPDE/test/runtests.jl | 20 ++ src/NeuralPDE.jl | 5 +- src/advancedHMC_MCMC.jl | 243 +----------------- src/pinn_types.jl | 34 --- 12 files changed, 312 insertions(+), 282 deletions(-) create mode 100644 lib/BayesianNeuralPDE/Project.toml rename {src => lib/BayesianNeuralPDE/src}/BPINN_ode.jl (99%) create mode 100644 lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl rename {src => lib/BayesianNeuralPDE/src}/PDE_BPINN.jl (99%) create mode 100644 lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl create mode 100644 lib/BayesianNeuralPDE/src/pinn_types.jl rename {test => lib/BayesianNeuralPDE/test}/BPINN_PDE_tests.jl (100%) rename {test => lib/BayesianNeuralPDE/test}/BPINN_tests.jl (100%) create mode 100644 lib/BayesianNeuralPDE/test/runtests.jl diff --git a/lib/BayesianNeuralPDE/Project.toml b/lib/BayesianNeuralPDE/Project.toml new file mode 100644 index 000000000..85a2fb49e --- /dev/null +++ b/lib/BayesianNeuralPDE/Project.toml @@ -0,0 +1,4 @@ +name = "BayesianNeuralPDE" +uuid = "3cea9122-e921-42ea-a9d7-c72fcb58ce53" +authors = ["paramthakkar123 "] +version = "0.1.0" \ No newline at end of file diff --git a/src/BPINN_ode.jl b/lib/BayesianNeuralPDE/src/BPINN_ode.jl similarity index 99% rename from src/BPINN_ode.jl rename to lib/BayesianNeuralPDE/src/BPINN_ode.jl index 489cd5f6d..95a34e9e4 100644 --- a/src/BPINN_ode.jl +++ b/lib/BayesianNeuralPDE/src/BPINN_ode.jl @@ -228,4 +228,4 @@ function SciMLBase.__solve(prob::SciMLBase.ODEProblem, alg::BNNODE, args...; dt end return BPINNsolution(fullsolution, ensemblecurves, estimnnparams, estimated_params, t) -end +end \ No newline at end of file diff --git a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl new file mode 100644 index 000000000..0bfbcb233 --- /dev/null +++ b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl @@ -0,0 +1,11 @@ +module BayesianNeuralPDE + +include("advancedHMC_MCMC.jl") +include("BPINN_ode.jl") +include("PDE_BPINN.jl") +include("pinn_types.jl") + +export BNNODE, ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde +export BPINNsolution, BayesianPINN + +end \ No newline at end of file diff --git a/src/PDE_BPINN.jl b/lib/BayesianNeuralPDE/src/PDE_BPINN.jl similarity index 99% rename from src/PDE_BPINN.jl rename to lib/BayesianNeuralPDE/src/PDE_BPINN.jl index 3ea00758f..b908a494e 100644 --- a/src/PDE_BPINN.jl +++ b/lib/BayesianNeuralPDE/src/PDE_BPINN.jl @@ -507,4 +507,4 @@ function ahmc_bayesian_pinn_pde(pde_system, discretization; return BPINNsolution( fullsolution, ensemblecurves, estimnnparams, estimated_params, timepoints) end -end +end \ No newline at end of file diff --git a/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl b/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl new file mode 100644 index 000000000..76810ba44 --- /dev/null +++ b/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl @@ -0,0 +1,240 @@ +""" + ahmc_bayesian_pinn_ode(prob, chain; strategy = GridTraining, dataset = [nothing], + init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0f0, + l2std = [0.05], phystd = [0.05], phynewstd = [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 + +```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) + +### 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] + +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 + solving equations +Incase you are only solving the Equations for solution, do not provide dataset + +## Positional Arguments + +* `prob`: DEProblem(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. +* `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`: standard deviation of new loss func term +* `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/ +* `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. +""" +function ahmc_bayesian_pinn_ode( + prob::SciMLBase.ODEProblem, chain; strategy = GridTraining, dataset = [nothing], + init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0, l2std = [0.05], + phystd = [0.05], phynewstd = [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, estim_collocate = false) + @assert !isinplace(prob) "The BPINN ODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t)." + + chain isa AbstractLuxLayer || (chain = FromFluxAdaptor()(chain)) + + strategy = strategy == GridTraining ? strategy(physdt) : strategy + + if dataset != [nothing] && + (length(dataset) < 2 || !(dataset isa Vector{<:Vector{<:AbstractFloat}})) + error("Invalid dataset. dataset would be timeseries (x̂,t) where type: Vector{Vector{AbstractFloat}") + end + + if dataset != [nothing] && param == [] + println("Dataset is only needed for Parameter Estimation + Forward Problem, not in only Forward Problem case.") + elseif dataset == [nothing] && param != [] + error("Dataset Required for Parameter Estimation.") + end + + initial_nnθ, chain, st = generate_ltd(chain, init_params) + + @assert nchains≤Threads.nthreads() "number of chains is greater than available threads" + @assert nchains≥1 "number of chains must be greater than 1" + + # eltype(physdt) cause needs Float64 for find_good_stepsize + # Lux chain(using component array later as vector_to_parameter need namedtuple) + T = eltype(physdt) + initial_θ = getdata(ComponentArray{T}(initial_nnθ)) + + # adding ode parameter estimation + nparameters = length(initial_θ) + ninv = length(param) + priors = [ + MvNormal(T(priorsNNw[1]) * ones(T, nparameters), + Diagonal(abs2.(T(priorsNNw[2]) .* ones(T, nparameters)))) + ] + + # append Ode params to all paramvector + if ninv > 0 + # shift ode params(initialise ode params by prior means) + initial_θ = vcat(initial_θ, [Distributions.params(param[i])[1] for i in 1:ninv]) + priors = vcat(priors, param) + nparameters += ninv + end + + smodel = StatefulLuxLayer{true}(chain, nothing, st) + # dimensions would be total no of params,initial_nnθ for Lux namedTuples + ℓπ = LogTargetDensity(nparameters, prob, smodel, strategy, dataset, priors, + phystd, phynewstd, l2std, autodiff, physdt, ninv, initial_nnθ, estim_collocate) + + if verbose + @printf("Current Physics Log-likelihood: %g\n", physloglikelihood(ℓπ, initial_θ)) + @printf("Current Prior Log-likelihood: %g\n", priorweights(ℓπ, initial_θ)) + @printf("Current SSE against dataset Log-likelihood: %g\n", + L2LossData(ℓπ, initial_θ)) + if estim_collocate + @printf("Current gradient loss against dataset Log-likelihood: %g\n", + L2loss2(ℓπ, initial_θ)) + end + end + + Adaptor = Adaptorkwargs[:Adaptor] + Metric = Adaptorkwargs[:Metric] + targetacceptancerate = Adaptorkwargs[:targetacceptancerate] + + # Define Hamiltonian system (nparameters ~ dimensionality of the sampling space) + metric = Metric(nparameters) + hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) + + # parallel sampling option + if nchains != 1 + # Cache to store the chains + chains = Vector{Any}(undef, nchains) + statsc = Vector{Any}(undef, nchains) + samplesc = Vector{Any}(undef, nchains) + + Threads.@threads for i in 1:nchains + # each chain has different initial NNparameter values(better posterior exploration) + initial_θ = vcat( + randn(eltype(initial_θ), nparameters - ninv), + initial_θ[(nparameters - ninv + 1):end] + ) + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = integratorchoice(Integratorkwargs, initial_ϵ) + adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), + StepSizeAdaptor(targetacceptancerate, integrator)) + + MCMC_alg = kernelchoice(Kernel, MCMCkwargs) + Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) + samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, adaptor; + progress = progress, verbose = verbose) + + samplesc[i] = samples + statsc[i] = stats + mcmc_chain = Chains(reduce(hcat, samples)') + chains[i] = mcmc_chain + end + + return chains, samplesc, statsc + else + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = integratorchoice(Integratorkwargs, initial_ϵ) + adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), + StepSizeAdaptor(targetacceptancerate, integrator)) + + MCMC_alg = kernelchoice(Kernel, MCMCkwargs) + Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) + samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, + adaptor; progress = progress, verbose = verbose) + + if verbose + println("Sampling Complete.") + @printf("Final Physics Log-likelihood: %g\n", + physloglikelihood(ℓπ, samples[end])) + @printf("Final Prior Log-likelihood: %g\n", priorweights(ℓπ, samples[end])) + @printf("Final SSE against dataset Log-likelihood: %g\n", + L2LossData(ℓπ, samples[end])) + if estim_collocate + @printf("Final gradient loss against dataset Log-likelihood: %g\n", + L2loss2(ℓπ, samples[end])) + end + end + + # return a chain(basic chain),samples and stats + matrix_samples = reshape(hcat(samples...), (length(samples[1]), length(samples), 1)) + mcmc_chain = MCMCChains.Chains(matrix_samples) + return mcmc_chain, samples, stats + end +end \ No newline at end of file diff --git a/lib/BayesianNeuralPDE/src/pinn_types.jl b/lib/BayesianNeuralPDE/src/pinn_types.jl new file mode 100644 index 000000000..65b365de5 --- /dev/null +++ b/lib/BayesianNeuralPDE/src/pinn_types.jl @@ -0,0 +1,33 @@ +""" + 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](https://turinglang.org/AdvancedHMC.jl/stable/) 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. +""" +@concrete struct BayesianPINN <: AbstractPINN + pinn <: PhysicsInformedNN + dataset +end + +function Base.getproperty(pinn::BayesianPINN, name::Symbol) + name === :dataset && return getfield(pinn, :dataset) + name === :pinn && return getfield(pinn, :pinn) + return getproperty(pinn.pinn, name) +end + +function BayesianPINN(args...; dataset = nothing, kwargs...) + dataset === nothing && (dataset = (nothing, nothing)) + return BayesianPINN(PhysicsInformedNN(args...; kwargs...), dataset) +end \ No newline at end of file diff --git a/test/BPINN_PDE_tests.jl b/lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl similarity index 100% rename from test/BPINN_PDE_tests.jl rename to lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl diff --git a/test/BPINN_tests.jl b/lib/BayesianNeuralPDE/test/BPINN_tests.jl similarity index 100% rename from test/BPINN_tests.jl rename to lib/BayesianNeuralPDE/test/BPINN_tests.jl diff --git a/lib/BayesianNeuralPDE/test/runtests.jl b/lib/BayesianNeuralPDE/test/runtests.jl new file mode 100644 index 000000000..4063f586e --- /dev/null +++ b/lib/BayesianNeuralPDE/test/runtests.jl @@ -0,0 +1,20 @@ +using ReTestItems, InteractiveUtils, Hwloc + +@info sprint(versioninfo) + +const GROUP = lowercase(get(ENV, "GROUP", "all")) + +const RETESTITEMS_NWORKERS = parse( + Int, get(ENV, "RETESTITEMS_NWORKERS", string(min(Hwloc.num_physical_cores(), 4)))) +const RETESTITEMS_NWORKER_THREADS = parse(Int, + get(ENV, "RETESTITEMS_NWORKER_THREADS", + string(max(Hwloc.num_virtual_cores() ÷ RETESTITEMS_NWORKERS, 1)))) + +using NeuralPDE + +@info "Running tests with $(RETESTITEMS_NWORKERS) workers and \ + $(RETESTITEMS_NWORKER_THREADS) threads for group $(GROUP)" + +ReTestItems.runtests(NeuralPDE; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]), + nworkers = RETESTITEMS_NWORKERS, + nworker_threads = RETESTITEMS_NWORKER_THREADS, testitem_timeout = 3600) \ No newline at end of file diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 23f2f18f2..771ec8841 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -83,19 +83,16 @@ include("dae_solve.jl") include("pino_ode_solve.jl") include("transform_inf_integral.jl") include("discretize.jl") +include("../lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl") include("neural_adapter.jl") include("advancedHMC_MCMC.jl") -include("BPINN_ode.jl") -include("PDE_BPINN.jl") include("dgm.jl") export PINOODE export NNODE, NNDAE -export BNNODE, ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde export PhysicsInformedNN, discretize -export BPINNsolution, BayesianPINN export DeepGalerkin export neural_adapter diff --git a/src/advancedHMC_MCMC.jl b/src/advancedHMC_MCMC.jl index f7f18e09b..726614255 100644 --- a/src/advancedHMC_MCMC.jl +++ b/src/advancedHMC_MCMC.jl @@ -259,245 +259,4 @@ function kernelchoice(Kernel, MCMCkwargs) else # HMC Kernel(MCMCkwargs[:n_leapfrog]) end -end - -""" - ahmc_bayesian_pinn_ode(prob, chain; strategy = GridTraining, dataset = [nothing], - init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0f0, - l2std = [0.05], phystd = [0.05], phynewstd = [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 - -```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) - -### 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] - -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 + solving equations -Incase you are only solving the Equations for solution, do not provide dataset - -## Positional Arguments - -* `prob`: DEProblem(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. -* `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`: standard deviation of new loss func term -* `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/ -* `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. -""" -function ahmc_bayesian_pinn_ode( - prob::SciMLBase.ODEProblem, chain; strategy = GridTraining, dataset = [nothing], - init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0, l2std = [0.05], - phystd = [0.05], phynewstd = [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, estim_collocate = false) - @assert !isinplace(prob) "The BPINN ODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t)." - - chain isa AbstractLuxLayer || (chain = FromFluxAdaptor()(chain)) - - strategy = strategy == GridTraining ? strategy(physdt) : strategy - - if dataset != [nothing] && - (length(dataset) < 2 || !(dataset isa Vector{<:Vector{<:AbstractFloat}})) - error("Invalid dataset. dataset would be timeseries (x̂,t) where type: Vector{Vector{AbstractFloat}") - end - - if dataset != [nothing] && param == [] - println("Dataset is only needed for Parameter Estimation + Forward Problem, not in only Forward Problem case.") - elseif dataset == [nothing] && param != [] - error("Dataset Required for Parameter Estimation.") - end - - initial_nnθ, chain, st = generate_ltd(chain, init_params) - - @assert nchains≤Threads.nthreads() "number of chains is greater than available threads" - @assert nchains≥1 "number of chains must be greater than 1" - - # eltype(physdt) cause needs Float64 for find_good_stepsize - # Lux chain(using component array later as vector_to_parameter need namedtuple) - T = eltype(physdt) - initial_θ = getdata(ComponentArray{T}(initial_nnθ)) - - # adding ode parameter estimation - nparameters = length(initial_θ) - ninv = length(param) - priors = [ - MvNormal(T(priorsNNw[1]) * ones(T, nparameters), - Diagonal(abs2.(T(priorsNNw[2]) .* ones(T, nparameters)))) - ] - - # append Ode params to all paramvector - if ninv > 0 - # shift ode params(initialise ode params by prior means) - initial_θ = vcat(initial_θ, [Distributions.params(param[i])[1] for i in 1:ninv]) - priors = vcat(priors, param) - nparameters += ninv - end - - smodel = StatefulLuxLayer{true}(chain, nothing, st) - # dimensions would be total no of params,initial_nnθ for Lux namedTuples - ℓπ = LogTargetDensity(nparameters, prob, smodel, strategy, dataset, priors, - phystd, phynewstd, l2std, autodiff, physdt, ninv, initial_nnθ, estim_collocate) - - if verbose - @printf("Current Physics Log-likelihood: %g\n", physloglikelihood(ℓπ, initial_θ)) - @printf("Current Prior Log-likelihood: %g\n", priorweights(ℓπ, initial_θ)) - @printf("Current SSE against dataset Log-likelihood: %g\n", - L2LossData(ℓπ, initial_θ)) - if estim_collocate - @printf("Current gradient loss against dataset Log-likelihood: %g\n", - L2loss2(ℓπ, initial_θ)) - end - end - - Adaptor = Adaptorkwargs[:Adaptor] - Metric = Adaptorkwargs[:Metric] - targetacceptancerate = Adaptorkwargs[:targetacceptancerate] - - # Define Hamiltonian system (nparameters ~ dimensionality of the sampling space) - metric = Metric(nparameters) - hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) - - # parallel sampling option - if nchains != 1 - # Cache to store the chains - chains = Vector{Any}(undef, nchains) - statsc = Vector{Any}(undef, nchains) - samplesc = Vector{Any}(undef, nchains) - - Threads.@threads for i in 1:nchains - # each chain has different initial NNparameter values(better posterior exploration) - initial_θ = vcat( - randn(eltype(initial_θ), nparameters - ninv), - initial_θ[(nparameters - ninv + 1):end] - ) - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = integratorchoice(Integratorkwargs, initial_ϵ) - adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), - StepSizeAdaptor(targetacceptancerate, integrator)) - - MCMC_alg = kernelchoice(Kernel, MCMCkwargs) - Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) - samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, adaptor; - progress = progress, verbose = verbose) - - samplesc[i] = samples - statsc[i] = stats - mcmc_chain = Chains(reduce(hcat, samples)') - chains[i] = mcmc_chain - end - - return chains, samplesc, statsc - else - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = integratorchoice(Integratorkwargs, initial_ϵ) - adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), - StepSizeAdaptor(targetacceptancerate, integrator)) - - MCMC_alg = kernelchoice(Kernel, MCMCkwargs) - Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) - samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, - adaptor; progress = progress, verbose = verbose) - - if verbose - println("Sampling Complete.") - @printf("Final Physics Log-likelihood: %g\n", - physloglikelihood(ℓπ, samples[end])) - @printf("Final Prior Log-likelihood: %g\n", priorweights(ℓπ, samples[end])) - @printf("Final SSE against dataset Log-likelihood: %g\n", - L2LossData(ℓπ, samples[end])) - if estim_collocate - @printf("Final gradient loss against dataset Log-likelihood: %g\n", - L2loss2(ℓπ, samples[end])) - end - end - - # return a chain(basic chain),samples and stats - matrix_samples = reshape(hcat(samples...), (length(samples[1]), length(samples), 1)) - mcmc_chain = MCMCChains.Chains(matrix_samples) - return mcmc_chain, samples, stats - end -end +end \ No newline at end of file diff --git a/src/pinn_types.jl b/src/pinn_types.jl index 6a7d617e9..33d992ba6 100644 --- a/src/pinn_types.jl +++ b/src/pinn_types.jl @@ -138,40 +138,6 @@ function PhysicsInformedNN( multioutput, kwargs) end -""" - 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](https://turinglang.org/AdvancedHMC.jl/stable/) 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. -""" -@concrete struct BayesianPINN <: AbstractPINN - pinn <: PhysicsInformedNN - dataset -end - -function Base.getproperty(pinn::BayesianPINN, name::Symbol) - name === :dataset && return getfield(pinn, :dataset) - name === :pinn && return getfield(pinn, :pinn) - return getproperty(pinn.pinn, name) -end - -function BayesianPINN(args...; dataset = nothing, kwargs...) - dataset === nothing && (dataset = (nothing, nothing)) - return BayesianPINN(PhysicsInformedNN(args...; kwargs...), dataset) -end - """ `PINNRepresentation`` From 75ac7c042257043978c381fc2a013a7abc867091 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Tue, 28 Jan 2025 07:59:04 +0530 Subject: [PATCH 02/15] Updates --- src/NeuralPDE.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 771ec8841..c24a59297 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -90,6 +90,10 @@ include("advancedHMC_MCMC.jl") include("dgm.jl") +using BayesianNeuralPDE +export BNNODE, ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde, + BPINNsolution, BayesianPINN + export PINOODE export NNODE, NNDAE export PhysicsInformedNN, discretize From 763dccca71d12ded50d5f1f7e0631fdc75a12496 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Tue, 28 Jan 2025 07:59:35 +0530 Subject: [PATCH 03/15] Updates --- src/NeuralPDE.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index c24a59297..d0593f8dd 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -83,7 +83,6 @@ include("dae_solve.jl") include("pino_ode_solve.jl") include("transform_inf_integral.jl") include("discretize.jl") -include("../lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl") include("neural_adapter.jl") include("advancedHMC_MCMC.jl") From 363b8052a0ee5bc23b96022612a8ea0b081d6494 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Tue, 28 Jan 2025 22:37:24 +0530 Subject: [PATCH 04/15] Circular --- src/NeuralPDE.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index d0593f8dd..22a18c7f8 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -90,8 +90,8 @@ include("advancedHMC_MCMC.jl") include("dgm.jl") using BayesianNeuralPDE -export BNNODE, ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde, - BPINNsolution, BayesianPINN +export ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde, + BNNODE, BPINNsolution, BayesianPINN export PINOODE export NNODE, NNDAE From a731c7c9967e2ba185cf6fb44013f00644375b13 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Tue, 28 Jan 2025 23:07:07 +0530 Subject: [PATCH 05/15] Updates --- .../src/BayesianNeuralPDE.jl | 1 + lib/BayesianNeuralPDE/src/discretize.jl | 78 ++++++++++++++++++ src/discretize.jl | 79 ------------------- 3 files changed, 79 insertions(+), 79 deletions(-) create mode 100644 lib/BayesianNeuralPDE/src/discretize.jl diff --git a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl index 0bfbcb233..a6e2f9d5e 100644 --- a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl +++ b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl @@ -2,6 +2,7 @@ module BayesianNeuralPDE include("advancedHMC_MCMC.jl") include("BPINN_ode.jl") +include("discretize.jl") include("PDE_BPINN.jl") include("pinn_types.jl") diff --git a/lib/BayesianNeuralPDE/src/discretize.jl b/lib/BayesianNeuralPDE/src/discretize.jl new file mode 100644 index 000000000..455ceadcb --- /dev/null +++ b/lib/BayesianNeuralPDE/src/discretize.jl @@ -0,0 +1,78 @@ +function get_likelihood_estimate_function(discretization::BayesianPINN) + dataset_pde, dataset_bc = discretization.dataset + + pde_loss_functions, bc_loss_functions = merge_strategy_with_loglikelihood_function( + pinnrep, strategy, + datafree_pde_loss_functions, datafree_bc_loss_functions) + + # required as Physics loss also needed on the discrete dataset domain points + # data points are discrete and so by default GridTraining loss applies + # passing placeholder dx with GridTraining, it uses data points irl + datapde_loss_functions, databc_loss_functions = if dataset_bc !== nothing || + dataset_pde !== nothing + merge_strategy_with_loglikelihood_function(pinnrep, GridTraining(0.1), + datafree_pde_loss_functions, datafree_bc_loss_functions, + train_sets_pde = dataset_pde, train_sets_bc = dataset_bc) + else + nothing, nothing + end + + # this includes losses from dataset domain points as well as discretization points + function full_loss_function(θ, allstd::Vector{Vector{Float64}}) + stdpdes, stdbcs, stdextra = allstd + # the aggregation happens on cpu even if the losses are gpu, probably fine since it's only a few of them + # SSE FOR LOSS ON GRIDPOINTS not MSE ! i, j depend on number of bcs and eqs + pde_loglikelihoods = sum([pde_loglike_function(θ, stdpdes[i]) + for (i, pde_loglike_function) in enumerate(pde_loss_functions)]) + + bc_loglikelihoods = sum([bc_loglike_function(θ, stdbcs[j]) + for (j, bc_loglike_function) in enumerate(bc_loss_functions)]) + + # final newloss creation components are similar to this + if !(datapde_loss_functions isa Nothing) + pde_loglikelihoods += sum([pde_loglike_function(θ, stdpdes[j]) + for (j, pde_loglike_function) in enumerate(datapde_loss_functions)]) + end + + if !(databc_loss_functions isa Nothing) + bc_loglikelihoods += sum([bc_loglike_function(θ, stdbcs[j]) + for (j, bc_loglike_function) in enumerate(databc_loss_functions)]) + end + + # this is kind of a hack, and means that whenever the outer function is evaluated the increment goes up, even if it's not being optimized + # that's why we prefer the user to maintain the increment in the outer loop callback during optimization + @ignore_derivatives if self_increment + iteration[] += 1 + end + + @ignore_derivatives begin + reweight_losses_func(θ, pde_loglikelihoods, + bc_loglikelihoods) + end + + weighted_pde_loglikelihood = adaloss.pde_loss_weights .* pde_loglikelihoods + weighted_bc_loglikelihood = adaloss.bc_loss_weights .* bc_loglikelihoods + + sum_weighted_pde_loglikelihood = sum(weighted_pde_loglikelihood) + sum_weighted_bc_loglikelihood = sum(weighted_bc_loglikelihood) + weighted_loglikelihood_before_additional = sum_weighted_pde_loglikelihood + + sum_weighted_bc_loglikelihood + + full_weighted_loglikelihood = if additional_loss isa Nothing + weighted_loglikelihood_before_additional + else + (θ_, p_) = param_estim ? (θ.depvar, θ.p) : (θ, nothing) + _additional_loss = additional_loss(phi, θ_, p_) + _additional_loglikelihood = logpdf(Normal(0, stdextra), _additional_loss) + + weighted_additional_loglikelihood = adaloss.additional_loss_weights[1] * + _additional_loglikelihood + + weighted_loglikelihood_before_additional + weighted_additional_loglikelihood + end + + return full_weighted_loglikelihood + end + + return full_loss_function +end \ No newline at end of file diff --git a/src/discretize.jl b/src/discretize.jl index 43653ba7c..e1778a3b9 100644 --- a/src/discretize.jl +++ b/src/discretize.jl @@ -523,85 +523,6 @@ function SciMLBase.symbolic_discretize(pde_system::PDESystem, discretization::Ab return full_loss_function end - function get_likelihood_estimate_function(discretization::BayesianPINN) - dataset_pde, dataset_bc = discretization.dataset - - pde_loss_functions, bc_loss_functions = merge_strategy_with_loglikelihood_function( - pinnrep, strategy, - datafree_pde_loss_functions, datafree_bc_loss_functions) - - # required as Physics loss also needed on the discrete dataset domain points - # data points are discrete and so by default GridTraining loss applies - # passing placeholder dx with GridTraining, it uses data points irl - datapde_loss_functions, databc_loss_functions = if dataset_bc !== nothing || - dataset_pde !== nothing - merge_strategy_with_loglikelihood_function(pinnrep, GridTraining(0.1), - datafree_pde_loss_functions, datafree_bc_loss_functions, - train_sets_pde = dataset_pde, train_sets_bc = dataset_bc) - else - nothing, nothing - end - - # this includes losses from dataset domain points as well as discretization points - function full_loss_function(θ, allstd::Vector{Vector{Float64}}) - stdpdes, stdbcs, stdextra = allstd - # the aggregation happens on cpu even if the losses are gpu, probably fine since it's only a few of them - # SSE FOR LOSS ON GRIDPOINTS not MSE ! i, j depend on number of bcs and eqs - pde_loglikelihoods = sum([pde_loglike_function(θ, stdpdes[i]) - for (i, pde_loglike_function) in enumerate(pde_loss_functions)]) - - bc_loglikelihoods = sum([bc_loglike_function(θ, stdbcs[j]) - for (j, bc_loglike_function) in enumerate(bc_loss_functions)]) - - # final newloss creation components are similar to this - if !(datapde_loss_functions isa Nothing) - pde_loglikelihoods += sum([pde_loglike_function(θ, stdpdes[j]) - for (j, pde_loglike_function) in enumerate(datapde_loss_functions)]) - end - - if !(databc_loss_functions isa Nothing) - bc_loglikelihoods += sum([bc_loglike_function(θ, stdbcs[j]) - for (j, bc_loglike_function) in enumerate(databc_loss_functions)]) - end - - # this is kind of a hack, and means that whenever the outer function is evaluated the increment goes up, even if it's not being optimized - # that's why we prefer the user to maintain the increment in the outer loop callback during optimization - @ignore_derivatives if self_increment - iteration[] += 1 - end - - @ignore_derivatives begin - reweight_losses_func(θ, pde_loglikelihoods, - bc_loglikelihoods) - end - - weighted_pde_loglikelihood = adaloss.pde_loss_weights .* pde_loglikelihoods - weighted_bc_loglikelihood = adaloss.bc_loss_weights .* bc_loglikelihoods - - sum_weighted_pde_loglikelihood = sum(weighted_pde_loglikelihood) - sum_weighted_bc_loglikelihood = sum(weighted_bc_loglikelihood) - weighted_loglikelihood_before_additional = sum_weighted_pde_loglikelihood + - sum_weighted_bc_loglikelihood - - full_weighted_loglikelihood = if additional_loss isa Nothing - weighted_loglikelihood_before_additional - else - (θ_, p_) = param_estim ? (θ.depvar, θ.p) : (θ, nothing) - _additional_loss = additional_loss(phi, θ_, p_) - _additional_loglikelihood = logpdf(Normal(0, stdextra), _additional_loss) - - weighted_additional_loglikelihood = adaloss.additional_loss_weights[1] * - _additional_loglikelihood - - weighted_loglikelihood_before_additional + weighted_additional_loglikelihood - end - - return full_weighted_loglikelihood - end - - return full_loss_function - end - full_loss_function = get_likelihood_estimate_function(discretization) pinnrep.loss_functions = PINNLossFunctions(bc_loss_functions, pde_loss_functions, full_loss_function, additional_loss, datafree_pde_loss_functions, From cbcc549e00113c5d3a60cb30851852e173414f6a Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Tue, 28 Jan 2025 23:45:27 +0530 Subject: [PATCH 06/15] Updates --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index b1e5bd7df..4948c2008 100644 --- a/Project.toml +++ b/Project.toml @@ -44,6 +44,7 @@ SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b" Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +BayesianNeuralPDE = "3cea9122-e921-42ea-a9d7-c72fcb58ce53" [weakdeps] TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f" From 8f9f71b74afbcd6630ff03f9b97c6bcfdc6f451c Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Wed, 29 Jan 2025 00:27:40 +0530 Subject: [PATCH 07/15] Updates --- Project.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/Project.toml b/Project.toml index 4948c2008..b1e5bd7df 100644 --- a/Project.toml +++ b/Project.toml @@ -44,7 +44,6 @@ SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b" Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" -BayesianNeuralPDE = "3cea9122-e921-42ea-a9d7-c72fcb58ce53" [weakdeps] TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f" From 5c1eb43f9d7c5900fc4d22aeb6c212ac432aecfb Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Thu, 30 Jan 2025 11:17:57 +0530 Subject: [PATCH 08/15] Fixes --- lib/BayesianNeuralPDE/test/runtests.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/lib/BayesianNeuralPDE/test/runtests.jl b/lib/BayesianNeuralPDE/test/runtests.jl index 4063f586e..d84f163db 100644 --- a/lib/BayesianNeuralPDE/test/runtests.jl +++ b/lib/BayesianNeuralPDE/test/runtests.jl @@ -4,11 +4,11 @@ using ReTestItems, InteractiveUtils, Hwloc const GROUP = lowercase(get(ENV, "GROUP", "all")) -const RETESTITEMS_NWORKERS = parse( - Int, get(ENV, "RETESTITEMS_NWORKERS", string(min(Hwloc.num_physical_cores(), 4)))) -const RETESTITEMS_NWORKER_THREADS = parse(Int, - get(ENV, "RETESTITEMS_NWORKER_THREADS", - string(max(Hwloc.num_virtual_cores() ÷ RETESTITEMS_NWORKERS, 1)))) +# const RETESTITEMS_NWORKERS = parse( +# Int, get(ENV, "RETESTITEMS_NWORKERS", string(min(Hwloc.num_physical_cores(), 4)))) +# const RETESTITEMS_NWORKER_THREADS = parse(Int, +# get(ENV, "RETESTITEMS_NWORKER_THREADS", +# string(max(Hwloc.num_virtual_cores() ÷ RETESTITEMS_NWORKERS, 1)))) using NeuralPDE From c4e61dfda9b37a7df75e844d3bd9aabaf32b6ad8 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Thu, 30 Jan 2025 11:22:49 +0530 Subject: [PATCH 09/15] Fixes --- lib/BayesianNeuralPDE/Project.toml | 14 +++++++++++++- lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl | 3 +++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/lib/BayesianNeuralPDE/Project.toml b/lib/BayesianNeuralPDE/Project.toml index 85a2fb49e..c316f8185 100644 --- a/lib/BayesianNeuralPDE/Project.toml +++ b/lib/BayesianNeuralPDE/Project.toml @@ -1,4 +1,16 @@ name = "BayesianNeuralPDE" uuid = "3cea9122-e921-42ea-a9d7-c72fcb58ce53" authors = ["paramthakkar123 "] -version = "0.1.0" \ No newline at end of file +version = "0.1.0" + +[deps] +AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" +ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1" +Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" \ No newline at end of file diff --git a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl index a6e2f9d5e..21af0d58d 100644 --- a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl +++ b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl @@ -1,5 +1,8 @@ module BayesianNeuralPDE +using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + include("advancedHMC_MCMC.jl") include("BPINN_ode.jl") include("discretize.jl") From cc96f2474268c0920e7a1a7bcfeece1fd20c7403 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sat, 1 Feb 2025 20:03:16 +0530 Subject: [PATCH 10/15] Updates --- lib/BayesianNeuralPDE/Project.toml | 20 ++++++++++++++++--- .../src/BayesianNeuralPDE.jl | 13 +++++++++++- src/NeuralPDE.jl | 8 ++++---- 3 files changed, 33 insertions(+), 8 deletions(-) diff --git a/lib/BayesianNeuralPDE/Project.toml b/lib/BayesianNeuralPDE/Project.toml index c316f8185..c6cc2b5c1 100644 --- a/lib/BayesianNeuralPDE/Project.toml +++ b/lib/BayesianNeuralPDE/Project.toml @@ -5,12 +5,26 @@ 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" -OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" -OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" 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" -Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" \ No newline at end of file +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" diff --git a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl index 21af0d58d..3944a8870 100644 --- a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl +++ b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl @@ -2,12 +2,23 @@ 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") -include("pinn_types.jl") export BNNODE, ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde export BPINNsolution, BayesianPINN diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 22a18c7f8..5810759d1 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -89,15 +89,15 @@ include("advancedHMC_MCMC.jl") include("dgm.jl") -using BayesianNeuralPDE -export ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde, - BNNODE, BPINNsolution, BayesianPINN - export PINOODE export NNODE, NNDAE export PhysicsInformedNN, discretize export DeepGalerkin +using BayesianNeuralPDE +export ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde, + BNNODE, BPINNsolution, BayesianPINN + export neural_adapter export GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, From acad60447031c77ac2b771e3594835be00ca3d0d Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sat, 1 Feb 2025 20:37:46 +0530 Subject: [PATCH 11/15] Fixes --- Project.toml | 1 + src/NeuralPDE.jl | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index b1e5bd7df..98c98c251 100644 --- a/Project.toml +++ b/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" diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 5810759d1..ba947fa7e 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -77,6 +77,7 @@ include("pinn_types.jl") include("symbolic_utilities.jl") include("training_strategies.jl") include("adaptive_losses.jl") +include("../lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl") include("ode_solve.jl") include("dae_solve.jl") @@ -94,7 +95,7 @@ export NNODE, NNDAE export PhysicsInformedNN, discretize export DeepGalerkin -using BayesianNeuralPDE +using .BayesianNeuralPDE export ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde, BNNODE, BPINNsolution, BayesianPINN From 82eae3120ba3d79ae0eed045b5b4c79ed3b79a42 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sun, 2 Feb 2025 09:37:10 +0530 Subject: [PATCH 12/15] Formatted each document and removed stale imports --- lib/BayesianNeuralPDE/src/BPINN_ode.jl | 272 ++--- .../src/BayesianNeuralPDE.jl | 2 +- lib/BayesianNeuralPDE/src/PDE_BPINN.jl | 840 ++++++++-------- lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl | 332 +++---- lib/BayesianNeuralPDE/src/discretize.jl | 124 +-- lib/BayesianNeuralPDE/src/pinn_types.jl | 18 +- lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl | 938 +++++++++--------- lib/BayesianNeuralPDE/test/BPINN_tests.jl | 834 ++++++++-------- lib/BayesianNeuralPDE/test/runtests.jl | 6 +- src/NeuralPDE.jl | 26 +- 10 files changed, 1695 insertions(+), 1697 deletions(-) diff --git a/lib/BayesianNeuralPDE/src/BPINN_ode.jl b/lib/BayesianNeuralPDE/src/BPINN_ode.jl index 95a34e9e4..68d996ba5 100644 --- a/lib/BayesianNeuralPDE/src/BPINN_ode.jl +++ b/lib/BayesianNeuralPDE/src/BPINN_ode.jl @@ -1,14 +1,14 @@ # 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) + 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 @@ -16,9 +16,9 @@ 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. + 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 @@ -48,14 +48,14 @@ 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) + 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) + priorsNNw = (0.0, 10.0), param = [Normal(6.5, 0.5), Normal(-3, 0.5)], + progress = true) sol_lux_pestim = solve(prob, alg) ``` @@ -78,42 +78,42 @@ Kevin Linka, Amelie Schäfer, Xuhui Meng, Zongren Zou, George Em Karniadakis, El "Bayesian Physics Informed Neural Networks for real-world nonlinear dynamical systems". """ @concrete struct BNNODE <: NeuralPDEAlgorithm - chain <: AbstractLuxLayer - kernel - 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 + 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) + 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 """ @@ -122,19 +122,19 @@ 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 + - n_steps + - acceptance_rate + - log_density + - hamiltonian_energy + - hamiltonian_energy_error + - numerical_error + - step_size + - nom_step_size """ @concrete struct BPINNstats - mcmc_chain - samples - statistics + mcmc_chain::Any + samples::Any + statistics::Any end """ @@ -148,84 +148,84 @@ contains fields related to that). parameters. """ @concrete struct BPINNsolution - original <: BPINNstats - ensemblesol - estimated_nn_params - estimated_de_params - timepoints + 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 \ No newline at end of file + 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 diff --git a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl index 3944a8870..38aaa696f 100644 --- a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl +++ b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl @@ -1,7 +1,7 @@ module BayesianNeuralPDE using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements using Printf: @printf using ConcreteStructs: @concrete using NeuralPDE: PhysicsInformedNN diff --git a/lib/BayesianNeuralPDE/src/PDE_BPINN.jl b/lib/BayesianNeuralPDE/src/PDE_BPINN.jl index b908a494e..beffe557e 100644 --- a/lib/BayesianNeuralPDE/src/PDE_BPINN.jl +++ b/lib/BayesianNeuralPDE/src/PDE_BPINN.jl @@ -1,263 +1,263 @@ @concrete struct PDELogTargetDensity - dim::Int - strategy <: AbstractTrainingStrategy - dataset <: Union{Nothing, Vector{<:Matrix{<:Real}}} - priors <: Vector{<:Distribution} - allstd::Vector{Vector{Float64}} - phynewstd::Vector{Float64} - names::Tuple - extraparams::Int - init_params <: Union{AbstractVector, NamedTuple, ComponentArray} - full_loglikelihood - L2_loss2 - Φ + dim::Int + strategy <: AbstractTrainingStrategy + dataset <: Union{Nothing, Vector{<:Matrix{<:Real}}} + priors <: Vector{<:Distribution} + allstd::Vector{Vector{Float64}} + phynewstd::Vector{Float64} + names::Tuple + extraparams::Int + init_params <: Union{AbstractVector, NamedTuple, ComponentArray} + full_loglikelihood::Any + L2_loss2::Any + Φ::Any end function LogDensityProblems.logdensity(ltd::PDELogTargetDensity, θ) - # for parameter estimation neccesarry to use multioutput case - if ltd.L2_loss2 === nothing - return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) + - priorlogpdf(ltd, θ) + L2LossData(ltd, θ) - else - return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) + - priorlogpdf(ltd, θ) + L2LossData(ltd, θ) + ltd.L2_loss2(setparameters(ltd, θ), ltd.phynewstd) - end + # for parameter estimation neccesarry to use multioutput case + if ltd.L2_loss2 === nothing + return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) + + priorlogpdf(ltd, θ) + L2LossData(ltd, θ) + else + return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) + + priorlogpdf(ltd, θ) + L2LossData(ltd, θ) + ltd.L2_loss2(setparameters(ltd, θ), ltd.phynewstd) + end end # you get a vector of losses function get_lossy(pinnrep, dataset, Dict_differentials) - eqs = pinnrep.eqs - depvars = pinnrep.depvars #depvar order is same as dataset - - # Dict_differentials is filled with Differential operator => diff_i key-value pairs - # masking operation - eqs_new = SymbolicUtils.substitute.(eqs, Ref(Dict_differentials)) - - to_subs, tobe_subs = get_symbols(dataset, depvars, eqs) - - # for values of all depvars at corresponding indvar values in dataset, create dictionaries {Dict(x(t) => 1.0496435863173237, y(t) => 1.9227770685615337)} - # In each Dict, num form of depvar is key to its value at certain coords of indvars, n_dicts = n_rows_dataset(or n_indvar_coords_dataset) - eq_subs = [Dict(tobe_subs[depvar] => to_subs[depvar][i] for depvar in depvars) - for i in 1:size(dataset[1][:, 1])[1]] - - # for each dataset point(eq_sub dictionary), substitute in masked equations - # n_collocated_equations = n_rows_dataset(or n_indvar_coords_dataset) - masked_colloc_equations = [[Symbolics.substitute(eq, eq_sub) for eq in eqs_new] - for eq_sub in eq_subs] - # now we have vector of dataset depvar's collocated equations - - # reverse dict for re-substituting values of Differential(t)(u(t)) etc - rev_Dict_differentials = Dict(value => key for (key, value) in Dict_differentials) - - # unmask Differential terms in masked_colloc_equations - colloc_equations = [Symbolics.substitute.( - masked_colloc_equation, Ref(rev_Dict_differentials)) - for masked_colloc_equation in masked_colloc_equations] - # nested vector of data_pde_loss_functions (as in discretize.jl) - # each sub vector has dataset's indvar coord's datafree_colloc_loss_function, n_subvectors = n_rows_dataset(or n_indvar_coords_dataset) - # zip each colloc equation with args for each build_loss call per equation vector - data_colloc_loss_functions = [[build_loss_function(pinnrep, eq, pde_indvar) - for (eq, pde_indvar, integration_indvar) in zip( - colloc_equation, - pinnrep.pde_indvars, - pinnrep.pde_integration_vars)] - for colloc_equation in colloc_equations] - - return data_colloc_loss_functions + eqs = pinnrep.eqs + depvars = pinnrep.depvars #depvar order is same as dataset + + # Dict_differentials is filled with Differential operator => diff_i key-value pairs + # masking operation + eqs_new = SymbolicUtils.substitute.(eqs, Ref(Dict_differentials)) + + to_subs, tobe_subs = get_symbols(dataset, depvars, eqs) + + # for values of all depvars at corresponding indvar values in dataset, create dictionaries {Dict(x(t) => 1.0496435863173237, y(t) => 1.9227770685615337)} + # In each Dict, num form of depvar is key to its value at certain coords of indvars, n_dicts = n_rows_dataset(or n_indvar_coords_dataset) + eq_subs = [Dict(tobe_subs[depvar] => to_subs[depvar][i] for depvar in depvars) + for i in 1:size(dataset[1][:, 1])[1]] + + # for each dataset point(eq_sub dictionary), substitute in masked equations + # n_collocated_equations = n_rows_dataset(or n_indvar_coords_dataset) + masked_colloc_equations = [[Symbolics.substitute(eq, eq_sub) for eq in eqs_new] + for eq_sub in eq_subs] + # now we have vector of dataset depvar's collocated equations + + # reverse dict for re-substituting values of Differential(t)(u(t)) etc + rev_Dict_differentials = Dict(value => key for (key, value) in Dict_differentials) + + # unmask Differential terms in masked_colloc_equations + colloc_equations = [Symbolics.substitute.( + masked_colloc_equation, Ref(rev_Dict_differentials)) + for masked_colloc_equation in masked_colloc_equations] + # nested vector of data_pde_loss_functions (as in discretize.jl) + # each sub vector has dataset's indvar coord's datafree_colloc_loss_function, n_subvectors = n_rows_dataset(or n_indvar_coords_dataset) + # zip each colloc equation with args for each build_loss call per equation vector + data_colloc_loss_functions = [[build_loss_function(pinnrep, eq, pde_indvar) + for (eq, pde_indvar, integration_indvar) in zip( + colloc_equation, + pinnrep.pde_indvars, + pinnrep.pde_integration_vars)] + for colloc_equation in colloc_equations] + + return data_colloc_loss_functions end function get_symbols(dataset, depvars, eqs) - # take only values of depvars from dataset - depvar_vals = [dataset_i[:, 1] for dataset_i in dataset] - # order of pinnrep.depvars, depvar_vals, BayesianPINN.dataset must be same - to_subs = Dict(depvars .=> depvar_vals) - - numform_vars = Symbolics.get_variables.(eqs) - Eq_vars = unique(reduce(vcat, numform_vars)) - # got equation's depvar num format {x(t)} for use in substitute() - - tobe_subs = Dict() - for a in depvars - for i in Eq_vars - expr = toexpr(i) - if (expr isa Expr) && (expr.args[1] == a) - tobe_subs[a] = i - end - end - end - # depvar symbolic and num format got, tobe_subs : Dict{Any, Any}(:y => y(t), :x => x(t)) - - return to_subs, tobe_subs + # take only values of depvars from dataset + depvar_vals = [dataset_i[:, 1] for dataset_i in dataset] + # order of pinnrep.depvars, depvar_vals, BayesianPINN.dataset must be same + to_subs = Dict(depvars .=> depvar_vals) + + numform_vars = Symbolics.get_variables.(eqs) + Eq_vars = unique(reduce(vcat, numform_vars)) + # got equation's depvar num format {x(t)} for use in substitute() + + tobe_subs = Dict() + for a in depvars + for i in Eq_vars + expr = toexpr(i) + if (expr isa Expr) && (expr.args[1] == a) + tobe_subs[a] = i + end + end + end + # depvar symbolic and num format got, tobe_subs : Dict{Any, Any}(:y => y(t), :x => x(t)) + + return to_subs, tobe_subs end @views function setparameters(ltd::PDELogTargetDensity, θ) - names = ltd.names - ps_new = θ[1:(end - ltd.extraparams)] - ps = ltd.init_params - - # multioutput case for Lux chains, for each depvar ps would contain Lux ComponentVectors - # which we use for mapping current ahmc sampled vector of parameters onto NNs - i = 0 - Luxparams = [vector_to_parameters(ps_new[((i += length(ps[x])) - length(ps[x]) + 1):i], - ps[x]) for x in names] - - a = ComponentArray(NamedTuple{ltd.names}(i for i in Luxparams)) - - if ltd.extraparams > 0 - return ComponentArray(; depvar = a, p = θ[(end - ltd.extraparams + 1):end]) - else - return ComponentArray(; depvar = a) - end + names = ltd.names + ps_new = θ[1:(end-ltd.extraparams)] + ps = ltd.init_params + + # multioutput case for Lux chains, for each depvar ps would contain Lux ComponentVectors + # which we use for mapping current ahmc sampled vector of parameters onto NNs + i = 0 + Luxparams = [vector_to_parameters(ps_new[((i+=length(ps[x]))-length(ps[x])+1):i], + ps[x]) for x in names] + + a = ComponentArray(NamedTuple{ltd.names}(i for i in Luxparams)) + + if ltd.extraparams > 0 + return ComponentArray(; depvar = a, p = θ[(end-ltd.extraparams+1):end]) + else + return ComponentArray(; depvar = a) + end end LogDensityProblems.dimension(ltd::PDELogTargetDensity) = ltd.dim function LogDensityProblems.capabilities(::PDELogTargetDensity) - LogDensityProblems.LogDensityOrder{1}() + LogDensityProblems.LogDensityOrder{1}() end # L2 losses loglikelihood(needed mainly for ODE parameter estimation) function L2LossData(ltd::PDELogTargetDensity, θ) - Φ = ltd.Φ - init_params = ltd.init_params - dataset = ltd.dataset - L2stds = ltd.allstd[3] - # each dep var has a diff dataset depending on its indep var and their domains - # these datasets are matrices of first col-dep var and remaining cols-all indep var - # ltd.init_params is needed to construct a vector of parameters into a ComponentVector - - # dataset of form Vector[matrix_x, matrix_y, matrix_z] - # matrix_i is of form [i,indvar1,indvar2,..] (needed in case if heterogenous domains) - - # Phi is the trial solution for each NN in chain array - # Creating logpdf( MvNormal(Phi(t,θ),std), dataset[i] ) - # dataset[i][:, 2:end] -> indepvar cols of a particular depvar's dataset - # dataset[i][:, 1] -> depvar col of depvar's dataset - - ltd.extraparams ≤ 0 && return false - - sumt = 0 - for i in eachindex(Φ) - sumt += logpdf( - MvNormal( - Φ[i](dataset[i][:, 2:end]', - vector_to_parameters(θ[1:(end - ltd.extraparams)], init_params)[ltd.names[i]])[ - 1, :], - Diagonal(abs2.(ones(size(dataset[i])[1]) .* L2stds[i]))), - dataset[i][:, 1]) - end - return sumt + Φ = ltd.Φ + init_params = ltd.init_params + dataset = ltd.dataset + L2stds = ltd.allstd[3] + # each dep var has a diff dataset depending on its indep var and their domains + # these datasets are matrices of first col-dep var and remaining cols-all indep var + # ltd.init_params is needed to construct a vector of parameters into a ComponentVector + + # dataset of form Vector[matrix_x, matrix_y, matrix_z] + # matrix_i is of form [i,indvar1,indvar2,..] (needed in case if heterogenous domains) + + # Phi is the trial solution for each NN in chain array + # Creating logpdf( MvNormal(Phi(t,θ),std), dataset[i] ) + # dataset[i][:, 2:end] -> indepvar cols of a particular depvar's dataset + # dataset[i][:, 1] -> depvar col of depvar's dataset + + ltd.extraparams ≤ 0 && return false + + sumt = 0 + for i in eachindex(Φ) + sumt += logpdf( + MvNormal( + Φ[i](dataset[i][:, 2:end]', + vector_to_parameters(θ[1:(end-ltd.extraparams)], init_params)[ltd.names[i]])[ + 1, :], + Diagonal(abs2.(ones(size(dataset[i])[1]) .* L2stds[i]))), + dataset[i][:, 1]) + end + return sumt end # priors for NN parameters + ODE constants function priorlogpdf(ltd::PDELogTargetDensity, θ) - allparams = ltd.priors - # Vector of ode parameters priors - invpriors = allparams[2:end] - nnwparams = allparams[1] + allparams = ltd.priors + # Vector of ode parameters priors + invpriors = allparams[2:end] + nnwparams = allparams[1] - ltd.extraparams ≤ 0 && return logpdf(nnwparams, θ) + ltd.extraparams ≤ 0 && return logpdf(nnwparams, θ) - invlogpdf = sum((length(θ) - ltd.extraparams + 1):length(θ)) do i - logpdf(invpriors[length(θ) - i + 1], θ[i]) - end + invlogpdf = sum((length(θ)-ltd.extraparams+1):length(θ)) do i + logpdf(invpriors[length(θ)-i+1], θ[i]) + end - return invlogpdf + logpdf(nnwparams, θ[1:(length(θ) - ltd.extraparams)]) + return invlogpdf + logpdf(nnwparams, θ[1:(length(θ)-ltd.extraparams)]) end function integratorchoice(Integratorkwargs, initial_ϵ) - Integrator = Integratorkwargs[:Integrator] - if Integrator == JitteredLeapfrog - jitter_rate = Integratorkwargs[:jitter_rate] - Integrator(initial_ϵ, jitter_rate) - elseif Integrator == TemperedLeapfrog - tempering_rate = Integratorkwargs[:tempering_rate] - Integrator(initial_ϵ, tempering_rate) - else - Integrator(initial_ϵ) - end + Integrator = Integratorkwargs[:Integrator] + if Integrator == JitteredLeapfrog + jitter_rate = Integratorkwargs[:jitter_rate] + Integrator(initial_ϵ, jitter_rate) + elseif Integrator == TemperedLeapfrog + tempering_rate = Integratorkwargs[:tempering_rate] + Integrator(initial_ϵ, tempering_rate) + else + Integrator(initial_ϵ) + end end function adaptorchoice(Adaptor, mma, ssa) - if Adaptor != AdvancedHMC.NoAdaptation() - Adaptor(mma, ssa) - else - AdvancedHMC.NoAdaptation() - end + if Adaptor != AdvancedHMC.NoAdaptation() + Adaptor(mma, ssa) + else + AdvancedHMC.NoAdaptation() + end end function inference(samples, pinnrep, saveats, numensemble, ℓπ) - domains = pinnrep.domains - phi = pinnrep.phi - dict_depvar_input = pinnrep.dict_depvar_input - depvars = pinnrep.depvars - - names = ℓπ.names - initial_nnθ = ℓπ.init_params - ninv = ℓπ.extraparams - - ranges = Dict([Symbol(d.variables) => infimum(d.domain):dx:supremum(d.domain) - for (d, dx) in zip(domains, saveats)]) - inputs = [dict_depvar_input[i] for i in depvars] - - span = [[ranges[indvar] for indvar in input] for input in inputs] - timepoints = [hcat(vec(map(points -> collect(points), - Iterators.product(span[i]...)))...) - for i in eachindex(phi)] - - # order of range's domains must match chain's inputs and dep_vars - samples = samples[(end - numensemble):end] - nnparams = length(samples[1][1:(end - ninv)]) - # get rows-ith param and col-ith sample value - estimnnparams = [Particles(reduce(hcat, samples)[i, :]) - for i in 1:nnparams] - - # PDE params - if ninv == 0 - estimated_params = [nothing] - else - estimated_params = [Particles(reduce(hcat, samples)[i, :]) - for i in (nnparams + 1):(nnparams + ninv)] - end - - # getting parameter ranges in case of Lux chains - Luxparams = [] - i = 0 - for x in names - len = length(initial_nnθ[x]) - push!(Luxparams, (i + 1):(i + len)) - i += len - end - - # convert to format directly usable by lux - estimatedLuxparams = [vector_to_parameters(estimnnparams[Luxparams[i]], - initial_nnθ[names[i]]) for i in eachindex(phi)] - - # infer predictions(preds) each row - NN, each col - ith sample - samplesn = reduce(hcat, samples) - preds = [] - for j in eachindex(phi) - push!(preds, - [phi[j](timepoints[j], - vector_to_parameters(samplesn[:, i][Luxparams[j]], - initial_nnθ[names[j]])) for i in 1:numensemble]) - end - - # note here no of samples referse to numensemble and points is the no of points in each dep_vars discretization - # each phi will give output in single domain of depvar(so we have each row as a vector of vector outputs) - # so we get after reduce a single matrix of n rows(samples), and j cols(points) - ensemblecurves = [Particles(reduce(vcat, preds[i])) for i in eachindex(phi)] - return ensemblecurves, estimatedLuxparams, estimated_params, timepoints + domains = pinnrep.domains + phi = pinnrep.phi + dict_depvar_input = pinnrep.dict_depvar_input + depvars = pinnrep.depvars + + names = ℓπ.names + initial_nnθ = ℓπ.init_params + ninv = ℓπ.extraparams + + ranges = Dict([Symbol(d.variables) => infimum(d.domain):dx:supremum(d.domain) + for (d, dx) in zip(domains, saveats)]) + inputs = [dict_depvar_input[i] for i in depvars] + + span = [[ranges[indvar] for indvar in input] for input in inputs] + timepoints = [hcat(vec(map(points -> collect(points), + Iterators.product(span[i]...)))...) + for i in eachindex(phi)] + + # order of range's domains must match chain's inputs and dep_vars + samples = samples[(end-numensemble):end] + nnparams = length(samples[1][1:(end-ninv)]) + # get rows-ith param and col-ith sample value + estimnnparams = [Particles(reduce(hcat, samples)[i, :]) + for i in 1:nnparams] + + # PDE params + if ninv == 0 + estimated_params = [nothing] + else + estimated_params = [Particles(reduce(hcat, samples)[i, :]) + for i in (nnparams+1):(nnparams+ninv)] + end + + # getting parameter ranges in case of Lux chains + Luxparams = [] + i = 0 + for x in names + len = length(initial_nnθ[x]) + push!(Luxparams, (i+1):(i+len)) + i += len + end + + # convert to format directly usable by lux + estimatedLuxparams = [vector_to_parameters(estimnnparams[Luxparams[i]], + initial_nnθ[names[i]]) for i in eachindex(phi)] + + # infer predictions(preds) each row - NN, each col - ith sample + samplesn = reduce(hcat, samples) + preds = [] + for j in eachindex(phi) + push!(preds, + [phi[j](timepoints[j], + vector_to_parameters(samplesn[:, i][Luxparams[j]], + initial_nnθ[names[j]])) for i in 1:numensemble]) + end + + # note here no of samples referse to numensemble and points is the no of points in each dep_vars discretization + # each phi will give output in single domain of depvar(so we have each row as a vector of vector outputs) + # so we get after reduce a single matrix of n rows(samples), and j cols(points) + ensemblecurves = [Particles(reduce(vcat, preds[i])) for i in eachindex(phi)] + return ensemblecurves, estimatedLuxparams, estimated_params, timepoints end """ - 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) + 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 @@ -304,207 +304,207 @@ end !!! warning - AdvancedHMC.jl is still developing convenience structs so might need changes on new - releases. + AdvancedHMC.jl is still developing convenience structs so might need changes on new + releases. """ function 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), Dict_differentials = nothing, progress = false, verbose = false) - pinnrep = symbolic_discretize(pde_system, discretization) - dataset_pde, dataset_bc = discretization.dataset - - newloss = if Dict_differentials isa Nothing - nothing - else - data_colloc_loss_functions = get_lossy(pinnrep, dataset_pde, Dict_differentials) - # size = number of indvar coords in dataset - # add case for if parameters present in bcs? - - train_sets_pde = get_dataset_train_points(pde_system.eqs, - dataset_pde, - pinnrep) - # j is number of indvar coords in dataset, i is number of PDE equations in system - # -1 is placeholder, removed in merge_strategy_with_loglikelihood_function function call (train_sets[:, 2:end]()) - colloc_train_sets = [[hcat([-1], train_sets_pde[i][:, j]...) - for i in eachindex(data_colloc_loss_functions[1])] - for j in eachindex(data_colloc_loss_functions)] - - # using dataset's indvar coords as train_sets_pde and indvar coord's datafree_colloc_loss_function, create loss functions - # placeholder strategy = GridTraining(0.1), datafree_bc_loss_function and train_sets_bc must be nothing - # order of indvar coords will be same as corresponding depvar coords values in dataset provided in get_lossy() call. - pde_loss_function_points = [merge_strategy_with_loglikelihood_function( - pinnrep, - GridTraining(0.1), - data_colloc_loss_functions[i], - nothing; - train_sets_pde = colloc_train_sets[i], - train_sets_bc = nothing)[1] - for i in eachindex(data_colloc_loss_functions)] - - function L2_loss2(θ, phynewstd) - # first sum is over points losses over many equations for the same points - # second sum is over all points - pde_loglikelihoods = sum([sum([pde_loss_function(θ, - phynewstd[i]) - for (i, pde_loss_function) in enumerate(pde_loss_functions)]) - for pde_loss_functions in pde_loss_function_points]) - end - end - - # add overall functionality for BC dataset points (case of parametric BC) ? - if ((dataset_bc isa Nothing) && (dataset_pde isa Nothing)) - dataset = nothing - elseif dataset_bc isa Nothing - dataset = dataset_pde - elseif dataset_pde isa Nothing - dataset = dataset_bc - else - dataset = [vcat(dataset_pde[i], dataset_bc[i]) for i in eachindex(dataset_pde)] - end - - if discretization.param_estim && isempty(param) - throw(UndefVarError(:param)) - elseif discretization.param_estim && dataset isa Nothing - throw(UndefVarError(:dataset)) - elseif discretization.param_estim && length(l2std) != length(pinnrep.depvars) - error("L2 stds length must match number of dependant variables") - end - - # for physics loglikelihood - full_weighted_loglikelihood = pinnrep.loss_functions.full_loss_function - chain = discretization.chain - - if length(pinnrep.domains) != length(saveats) - error("Number of independent variables must match saveat inference discretization steps") - end - - # NN solutions for loglikelihood which is used for L2lossdata - Φ = pinnrep.phi - - @assert nchains≥1 "number of chains must be greater than or equal to 1" - - # remove inv params take only NN params, AHMC uses Float64 - initial_nnθ = pinnrep.flat_init_params[1:(end - length(param))] - initial_θ = collect(Float64, initial_nnθ) - - # contains only NN parameters - initial_nnθ = pinnrep.init_params - - names = ntuple(i -> pinnrep.depvars[i], length(chain)) - - #ode parameter estimation - nparameters = length(initial_θ) - ninv = length(param) - # add init_params for NN params - priors = [ - MvNormal(priorsNNw[1] * ones(nparameters), - Diagonal(abs2.(priorsNNw[2] .* ones(nparameters)))) - ] - - # append Ode params to all paramvector - initial_θ - if ninv > 0 - # shift ode params(initialise ode params by prior means) - # check if means or user specified is better - initial_θ = vcat(initial_θ, [Distributions.params(param[i])[1] for i in 1:ninv]) - priors = vcat(priors, param) - nparameters += ninv - end - - # vector in case of N-dimensional domains - strategy = discretization.strategy - - # dimensions would be total no of params,initial_nnθ for Lux namedTuples - ℓπ = PDELogTargetDensity( - nparameters, strategy, dataset, priors, [phystd, bcstd, l2std], phynewstd, - names, ninv, initial_nnθ, full_weighted_loglikelihood, newloss, Φ) - - Adaptor, Metric, targetacceptancerate = Adaptorkwargs[:Adaptor], - Adaptorkwargs[:Metric], Adaptorkwargs[:targetacceptancerate] - - # Define Hamiltonian system (nparameters ~ dimensionality of the sampling space) - metric = Metric(nparameters) - hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) - - if verbose - @printf("Current Physics Log-likelihood : %g\n", - ℓπ.full_loglikelihood(setparameters(ℓπ, initial_θ), ℓπ.allstd)) - @printf("Current Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, initial_θ)) - @printf("Current SSE against dataset Log-likelihood : %g\n", - L2LossData(ℓπ, initial_θ)) - if !(newloss isa Nothing) - @printf("Current new loss : %g\n", - ℓπ.L2_loss2(setparameters(ℓπ, initial_θ), - ℓπ.phynewstd)) - end - end - - # parallel sampling option - if nchains != 1 - - # Cache to store the chains - bpinnsols = Vector{Any}(undef, nchains) - - Threads.@threads for i in 1:nchains - # each chain has different initial NNparameter values(better posterior exploration) - initial_θ = vcat(randn(nparameters - ninv), - initial_θ[(nparameters - ninv + 1):end]) - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = integratorchoice(Integratorkwargs, initial_ϵ) - adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), - StepSizeAdaptor(targetacceptancerate, integrator)) - Kernel = AdvancedHMC.make_kernel(Kernel, integrator) - samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, adaptor; - progress = progress, verbose = verbose) - - # return a chain(basic chain),samples and stats - matrix_samples = hcat(samples...) - mcmc_chain = MCMCChains.Chains(matrix_samples') - - fullsolution = BPINNstats(mcmc_chain, samples, stats) - ensemblecurves, estimnnparams, estimated_params, timepoints = inference( - samples, pinnrep, saveat, numensemble, ℓπ) - - bpinnsols[i] = BPINNsolution( - fullsolution, ensemblecurves, estimnnparams, estimated_params, timepoints) - end - return bpinnsols - else - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = integratorchoice(Integratorkwargs, initial_ϵ) - adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), - StepSizeAdaptor(targetacceptancerate, integrator)) - - Kernel = AdvancedHMC.make_kernel(Kernel, integrator) - samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, - adaptor; progress = progress, verbose = verbose) - - # return a chain(basic chain),samples and stats - matrix_samples = hcat(samples...) - mcmc_chain = MCMCChains.Chains(matrix_samples') - - if verbose - @printf("Sampling Complete.\n") - @printf("Final Physics Log-likelihood : %g\n", - ℓπ.full_loglikelihood(setparameters(ℓπ, samples[end]), ℓπ.allstd)) - @printf("Final Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, samples[end])) - @printf("Final SSE against dataset Log-likelihood : %g\n", - L2LossData(ℓπ, samples[end])) - if !(newloss isa Nothing) - @printf("Final new loss : %g\n", - ℓπ.L2_loss2(setparameters(ℓπ, samples[end]), - ℓπ.phynewstd)) - end - end - - fullsolution = BPINNstats(mcmc_chain, samples, stats) - ensemblecurves, estimnnparams, estimated_params, timepoints = inference(samples, - pinnrep, saveats, numensemble, ℓπ) - - return BPINNsolution( - fullsolution, ensemblecurves, estimnnparams, estimated_params, timepoints) - end -end \ No newline at end of file + 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), Dict_differentials = nothing, progress = false, verbose = false) + pinnrep = symbolic_discretize(pde_system, discretization) + dataset_pde, dataset_bc = discretization.dataset + + newloss = if Dict_differentials isa Nothing + nothing + else + data_colloc_loss_functions = get_lossy(pinnrep, dataset_pde, Dict_differentials) + # size = number of indvar coords in dataset + # add case for if parameters present in bcs? + + train_sets_pde = get_dataset_train_points(pde_system.eqs, + dataset_pde, + pinnrep) + # j is number of indvar coords in dataset, i is number of PDE equations in system + # -1 is placeholder, removed in merge_strategy_with_loglikelihood_function function call (train_sets[:, 2:end]()) + colloc_train_sets = [[hcat([-1], train_sets_pde[i][:, j]...) + for i in eachindex(data_colloc_loss_functions[1])] + for j in eachindex(data_colloc_loss_functions)] + + # using dataset's indvar coords as train_sets_pde and indvar coord's datafree_colloc_loss_function, create loss functions + # placeholder strategy = GridTraining(0.1), datafree_bc_loss_function and train_sets_bc must be nothing + # order of indvar coords will be same as corresponding depvar coords values in dataset provided in get_lossy() call. + pde_loss_function_points = [merge_strategy_with_loglikelihood_function( + pinnrep, + GridTraining(0.1), + data_colloc_loss_functions[i], + nothing; + train_sets_pde = colloc_train_sets[i], + train_sets_bc = nothing)[1] + for i in eachindex(data_colloc_loss_functions)] + + function L2_loss2(θ, phynewstd) + # first sum is over points losses over many equations for the same points + # second sum is over all points + pde_loglikelihoods = sum([sum([pde_loss_function(θ, + phynewstd[i]) + for (i, pde_loss_function) in enumerate(pde_loss_functions)]) + for pde_loss_functions in pde_loss_function_points]) + end + end + + # add overall functionality for BC dataset points (case of parametric BC) ? + if ((dataset_bc isa Nothing) && (dataset_pde isa Nothing)) + dataset = nothing + elseif dataset_bc isa Nothing + dataset = dataset_pde + elseif dataset_pde isa Nothing + dataset = dataset_bc + else + dataset = [vcat(dataset_pde[i], dataset_bc[i]) for i in eachindex(dataset_pde)] + end + + if discretization.param_estim && isempty(param) + throw(UndefVarError(:param)) + elseif discretization.param_estim && dataset isa Nothing + throw(UndefVarError(:dataset)) + elseif discretization.param_estim && length(l2std) != length(pinnrep.depvars) + error("L2 stds length must match number of dependant variables") + end + + # for physics loglikelihood + full_weighted_loglikelihood = pinnrep.loss_functions.full_loss_function + chain = discretization.chain + + if length(pinnrep.domains) != length(saveats) + error("Number of independent variables must match saveat inference discretization steps") + end + + # NN solutions for loglikelihood which is used for L2lossdata + Φ = pinnrep.phi + + @assert nchains ≥ 1 "number of chains must be greater than or equal to 1" + + # remove inv params take only NN params, AHMC uses Float64 + initial_nnθ = pinnrep.flat_init_params[1:(end-length(param))] + initial_θ = collect(Float64, initial_nnθ) + + # contains only NN parameters + initial_nnθ = pinnrep.init_params + + names = ntuple(i -> pinnrep.depvars[i], length(chain)) + + #ode parameter estimation + nparameters = length(initial_θ) + ninv = length(param) + # add init_params for NN params + priors = [ + MvNormal(priorsNNw[1] * ones(nparameters), + Diagonal(abs2.(priorsNNw[2] .* ones(nparameters)))), + ] + + # append Ode params to all paramvector - initial_θ + if ninv > 0 + # shift ode params(initialise ode params by prior means) + # check if means or user specified is better + initial_θ = vcat(initial_θ, [Distributions.params(param[i])[1] for i in 1:ninv]) + priors = vcat(priors, param) + nparameters += ninv + end + + # vector in case of N-dimensional domains + strategy = discretization.strategy + + # dimensions would be total no of params,initial_nnθ for Lux namedTuples + ℓπ = PDELogTargetDensity( + nparameters, strategy, dataset, priors, [phystd, bcstd, l2std], phynewstd, + names, ninv, initial_nnθ, full_weighted_loglikelihood, newloss, Φ) + + Adaptor, Metric, targetacceptancerate = Adaptorkwargs[:Adaptor], + Adaptorkwargs[:Metric], Adaptorkwargs[:targetacceptancerate] + + # Define Hamiltonian system (nparameters ~ dimensionality of the sampling space) + metric = Metric(nparameters) + hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) + + if verbose + @printf("Current Physics Log-likelihood : %g\n", + ℓπ.full_loglikelihood(setparameters(ℓπ, initial_θ), ℓπ.allstd)) + @printf("Current Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, initial_θ)) + @printf("Current SSE against dataset Log-likelihood : %g\n", + L2LossData(ℓπ, initial_θ)) + if !(newloss isa Nothing) + @printf("Current new loss : %g\n", + ℓπ.L2_loss2(setparameters(ℓπ, initial_θ), + ℓπ.phynewstd)) + end + end + + # parallel sampling option + if nchains != 1 + + # Cache to store the chains + bpinnsols = Vector{Any}(undef, nchains) + + Threads.@threads for i in 1:nchains + # each chain has different initial NNparameter values(better posterior exploration) + initial_θ = vcat(randn(nparameters - ninv), + initial_θ[(nparameters-ninv+1):end]) + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = integratorchoice(Integratorkwargs, initial_ϵ) + adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), + StepSizeAdaptor(targetacceptancerate, integrator)) + Kernel = AdvancedHMC.make_kernel(Kernel, integrator) + samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, adaptor; + progress = progress, verbose = verbose) + + # return a chain(basic chain),samples and stats + matrix_samples = hcat(samples...) + mcmc_chain = MCMCChains.Chains(matrix_samples') + + fullsolution = BPINNstats(mcmc_chain, samples, stats) + ensemblecurves, estimnnparams, estimated_params, timepoints = inference( + samples, pinnrep, saveat, numensemble, ℓπ) + + bpinnsols[i] = BPINNsolution( + fullsolution, ensemblecurves, estimnnparams, estimated_params, timepoints) + end + return bpinnsols + else + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = integratorchoice(Integratorkwargs, initial_ϵ) + adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), + StepSizeAdaptor(targetacceptancerate, integrator)) + + Kernel = AdvancedHMC.make_kernel(Kernel, integrator) + samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, + adaptor; progress = progress, verbose = verbose) + + # return a chain(basic chain),samples and stats + matrix_samples = hcat(samples...) + mcmc_chain = MCMCChains.Chains(matrix_samples') + + if verbose + @printf("Sampling Complete.\n") + @printf("Final Physics Log-likelihood : %g\n", + ℓπ.full_loglikelihood(setparameters(ℓπ, samples[end]), ℓπ.allstd)) + @printf("Final Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, samples[end])) + @printf("Final SSE against dataset Log-likelihood : %g\n", + L2LossData(ℓπ, samples[end])) + if !(newloss isa Nothing) + @printf("Final new loss : %g\n", + ℓπ.L2_loss2(setparameters(ℓπ, samples[end]), + ℓπ.phynewstd)) + end + end + + fullsolution = BPINNstats(mcmc_chain, samples, stats) + ensemblecurves, estimnnparams, estimated_params, timepoints = inference(samples, + pinnrep, saveats, numensemble, ℓπ) + + return BPINNsolution( + fullsolution, ensemblecurves, estimnnparams, estimated_params, timepoints) + end +end diff --git a/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl b/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl index 76810ba44..ff0c732da 100644 --- a/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl +++ b/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl @@ -1,19 +1,19 @@ """ - ahmc_bayesian_pinn_ode(prob, chain; strategy = GridTraining, dataset = [nothing], - init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0f0, - l2std = [0.05], phystd = [0.05], phynewstd = [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) + ahmc_bayesian_pinn_ode(prob, chain; strategy = GridTraining, dataset = [nothing], + init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0f0, + l2std = [0.05], phystd = [0.05], phynewstd = [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. + 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 @@ -37,20 +37,20 @@ 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)) + 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)]) + 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 @@ -88,153 +88,153 @@ Incase you are only solving the Equations for solution, do not provide dataset 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/ + * `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/ * `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. + AdvancedHMC.jl is still developing convenience structs so might need changes on new + releases. """ function ahmc_bayesian_pinn_ode( - prob::SciMLBase.ODEProblem, chain; strategy = GridTraining, dataset = [nothing], - init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0, l2std = [0.05], - phystd = [0.05], phynewstd = [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, estim_collocate = false) - @assert !isinplace(prob) "The BPINN ODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t)." - - chain isa AbstractLuxLayer || (chain = FromFluxAdaptor()(chain)) - - strategy = strategy == GridTraining ? strategy(physdt) : strategy - - if dataset != [nothing] && - (length(dataset) < 2 || !(dataset isa Vector{<:Vector{<:AbstractFloat}})) - error("Invalid dataset. dataset would be timeseries (x̂,t) where type: Vector{Vector{AbstractFloat}") - end - - if dataset != [nothing] && param == [] - println("Dataset is only needed for Parameter Estimation + Forward Problem, not in only Forward Problem case.") - elseif dataset == [nothing] && param != [] - error("Dataset Required for Parameter Estimation.") - end - - initial_nnθ, chain, st = generate_ltd(chain, init_params) - - @assert nchains≤Threads.nthreads() "number of chains is greater than available threads" - @assert nchains≥1 "number of chains must be greater than 1" - - # eltype(physdt) cause needs Float64 for find_good_stepsize - # Lux chain(using component array later as vector_to_parameter need namedtuple) - T = eltype(physdt) - initial_θ = getdata(ComponentArray{T}(initial_nnθ)) - - # adding ode parameter estimation - nparameters = length(initial_θ) - ninv = length(param) - priors = [ - MvNormal(T(priorsNNw[1]) * ones(T, nparameters), - Diagonal(abs2.(T(priorsNNw[2]) .* ones(T, nparameters)))) - ] - - # append Ode params to all paramvector - if ninv > 0 - # shift ode params(initialise ode params by prior means) - initial_θ = vcat(initial_θ, [Distributions.params(param[i])[1] for i in 1:ninv]) - priors = vcat(priors, param) - nparameters += ninv - end - - smodel = StatefulLuxLayer{true}(chain, nothing, st) - # dimensions would be total no of params,initial_nnθ for Lux namedTuples - ℓπ = LogTargetDensity(nparameters, prob, smodel, strategy, dataset, priors, - phystd, phynewstd, l2std, autodiff, physdt, ninv, initial_nnθ, estim_collocate) - - if verbose - @printf("Current Physics Log-likelihood: %g\n", physloglikelihood(ℓπ, initial_θ)) - @printf("Current Prior Log-likelihood: %g\n", priorweights(ℓπ, initial_θ)) - @printf("Current SSE against dataset Log-likelihood: %g\n", - L2LossData(ℓπ, initial_θ)) - if estim_collocate - @printf("Current gradient loss against dataset Log-likelihood: %g\n", - L2loss2(ℓπ, initial_θ)) - end - end - - Adaptor = Adaptorkwargs[:Adaptor] - Metric = Adaptorkwargs[:Metric] - targetacceptancerate = Adaptorkwargs[:targetacceptancerate] - - # Define Hamiltonian system (nparameters ~ dimensionality of the sampling space) - metric = Metric(nparameters) - hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) - - # parallel sampling option - if nchains != 1 - # Cache to store the chains - chains = Vector{Any}(undef, nchains) - statsc = Vector{Any}(undef, nchains) - samplesc = Vector{Any}(undef, nchains) - - Threads.@threads for i in 1:nchains - # each chain has different initial NNparameter values(better posterior exploration) - initial_θ = vcat( - randn(eltype(initial_θ), nparameters - ninv), - initial_θ[(nparameters - ninv + 1):end] - ) - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = integratorchoice(Integratorkwargs, initial_ϵ) - adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), - StepSizeAdaptor(targetacceptancerate, integrator)) - - MCMC_alg = kernelchoice(Kernel, MCMCkwargs) - Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) - samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, adaptor; - progress = progress, verbose = verbose) - - samplesc[i] = samples - statsc[i] = stats - mcmc_chain = Chains(reduce(hcat, samples)') - chains[i] = mcmc_chain - end - - return chains, samplesc, statsc - else - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = integratorchoice(Integratorkwargs, initial_ϵ) - adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), - StepSizeAdaptor(targetacceptancerate, integrator)) - - MCMC_alg = kernelchoice(Kernel, MCMCkwargs) - Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) - samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, - adaptor; progress = progress, verbose = verbose) - - if verbose - println("Sampling Complete.") - @printf("Final Physics Log-likelihood: %g\n", - physloglikelihood(ℓπ, samples[end])) - @printf("Final Prior Log-likelihood: %g\n", priorweights(ℓπ, samples[end])) - @printf("Final SSE against dataset Log-likelihood: %g\n", - L2LossData(ℓπ, samples[end])) - if estim_collocate - @printf("Final gradient loss against dataset Log-likelihood: %g\n", - L2loss2(ℓπ, samples[end])) - end - end - - # return a chain(basic chain),samples and stats - matrix_samples = reshape(hcat(samples...), (length(samples[1]), length(samples), 1)) - mcmc_chain = MCMCChains.Chains(matrix_samples) - return mcmc_chain, samples, stats - end -end \ No newline at end of file + prob::SciMLBase.ODEProblem, chain; strategy = GridTraining, dataset = [nothing], + init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0, l2std = [0.05], + phystd = [0.05], phynewstd = [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, estim_collocate = false) + @assert !isinplace(prob) "The BPINN ODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t)." + + chain isa AbstractLuxLayer || (chain = FromFluxAdaptor()(chain)) + + strategy = strategy == GridTraining ? strategy(physdt) : strategy + + if dataset != [nothing] && + (length(dataset) < 2 || !(dataset isa Vector{<:Vector{<:AbstractFloat}})) + error("Invalid dataset. dataset would be timeseries (x̂,t) where type: Vector{Vector{AbstractFloat}") + end + + if dataset != [nothing] && param == [] + println("Dataset is only needed for Parameter Estimation + Forward Problem, not in only Forward Problem case.") + elseif dataset == [nothing] && param != [] + error("Dataset Required for Parameter Estimation.") + end + + initial_nnθ, chain, st = generate_ltd(chain, init_params) + + @assert nchains ≤ Threads.nthreads() "number of chains is greater than available threads" + @assert nchains ≥ 1 "number of chains must be greater than 1" + + # eltype(physdt) cause needs Float64 for find_good_stepsize + # Lux chain(using component array later as vector_to_parameter need namedtuple) + T = eltype(physdt) + initial_θ = getdata(ComponentArray{T}(initial_nnθ)) + + # adding ode parameter estimation + nparameters = length(initial_θ) + ninv = length(param) + priors = [ + MvNormal(T(priorsNNw[1]) * ones(T, nparameters), + Diagonal(abs2.(T(priorsNNw[2]) .* ones(T, nparameters)))), + ] + + # append Ode params to all paramvector + if ninv > 0 + # shift ode params(initialise ode params by prior means) + initial_θ = vcat(initial_θ, [Distributions.params(param[i])[1] for i in 1:ninv]) + priors = vcat(priors, param) + nparameters += ninv + end + + smodel = StatefulLuxLayer{true}(chain, nothing, st) + # dimensions would be total no of params,initial_nnθ for Lux namedTuples + ℓπ = LogTargetDensity(nparameters, prob, smodel, strategy, dataset, priors, + phystd, phynewstd, l2std, autodiff, physdt, ninv, initial_nnθ, estim_collocate) + + if verbose + @printf("Current Physics Log-likelihood: %g\n", physloglikelihood(ℓπ, initial_θ)) + @printf("Current Prior Log-likelihood: %g\n", priorweights(ℓπ, initial_θ)) + @printf("Current SSE against dataset Log-likelihood: %g\n", + L2LossData(ℓπ, initial_θ)) + if estim_collocate + @printf("Current gradient loss against dataset Log-likelihood: %g\n", + L2loss2(ℓπ, initial_θ)) + end + end + + Adaptor = Adaptorkwargs[:Adaptor] + Metric = Adaptorkwargs[:Metric] + targetacceptancerate = Adaptorkwargs[:targetacceptancerate] + + # Define Hamiltonian system (nparameters ~ dimensionality of the sampling space) + metric = Metric(nparameters) + hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) + + # parallel sampling option + if nchains != 1 + # Cache to store the chains + chains = Vector{Any}(undef, nchains) + statsc = Vector{Any}(undef, nchains) + samplesc = Vector{Any}(undef, nchains) + + Threads.@threads for i in 1:nchains + # each chain has different initial NNparameter values(better posterior exploration) + initial_θ = vcat( + randn(eltype(initial_θ), nparameters - ninv), + initial_θ[(nparameters-ninv+1):end], + ) + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = integratorchoice(Integratorkwargs, initial_ϵ) + adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), + StepSizeAdaptor(targetacceptancerate, integrator)) + + MCMC_alg = kernelchoice(Kernel, MCMCkwargs) + Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) + samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, adaptor; + progress = progress, verbose = verbose) + + samplesc[i] = samples + statsc[i] = stats + mcmc_chain = Chains(reduce(hcat, samples)') + chains[i] = mcmc_chain + end + + return chains, samplesc, statsc + else + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = integratorchoice(Integratorkwargs, initial_ϵ) + adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), + StepSizeAdaptor(targetacceptancerate, integrator)) + + MCMC_alg = kernelchoice(Kernel, MCMCkwargs) + Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) + samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, + adaptor; progress = progress, verbose = verbose) + + if verbose + println("Sampling Complete.") + @printf("Final Physics Log-likelihood: %g\n", + physloglikelihood(ℓπ, samples[end])) + @printf("Final Prior Log-likelihood: %g\n", priorweights(ℓπ, samples[end])) + @printf("Final SSE against dataset Log-likelihood: %g\n", + L2LossData(ℓπ, samples[end])) + if estim_collocate + @printf("Final gradient loss against dataset Log-likelihood: %g\n", + L2loss2(ℓπ, samples[end])) + end + end + + # return a chain(basic chain),samples and stats + matrix_samples = reshape(hcat(samples...), (length(samples[1]), length(samples), 1)) + mcmc_chain = MCMCChains.Chains(matrix_samples) + return mcmc_chain, samples, stats + end +end diff --git a/lib/BayesianNeuralPDE/src/discretize.jl b/lib/BayesianNeuralPDE/src/discretize.jl index 455ceadcb..e7fa68941 100644 --- a/lib/BayesianNeuralPDE/src/discretize.jl +++ b/lib/BayesianNeuralPDE/src/discretize.jl @@ -1,78 +1,78 @@ function get_likelihood_estimate_function(discretization::BayesianPINN) - dataset_pde, dataset_bc = discretization.dataset + dataset_pde, dataset_bc = discretization.dataset - pde_loss_functions, bc_loss_functions = merge_strategy_with_loglikelihood_function( - pinnrep, strategy, - datafree_pde_loss_functions, datafree_bc_loss_functions) + pde_loss_functions, bc_loss_functions = merge_strategy_with_loglikelihood_function( + pinnrep, strategy, + datafree_pde_loss_functions, datafree_bc_loss_functions) - # required as Physics loss also needed on the discrete dataset domain points - # data points are discrete and so by default GridTraining loss applies - # passing placeholder dx with GridTraining, it uses data points irl - datapde_loss_functions, databc_loss_functions = if dataset_bc !== nothing || - dataset_pde !== nothing - merge_strategy_with_loglikelihood_function(pinnrep, GridTraining(0.1), - datafree_pde_loss_functions, datafree_bc_loss_functions, - train_sets_pde = dataset_pde, train_sets_bc = dataset_bc) - else - nothing, nothing - end + # required as Physics loss also needed on the discrete dataset domain points + # data points are discrete and so by default GridTraining loss applies + # passing placeholder dx with GridTraining, it uses data points irl + datapde_loss_functions, databc_loss_functions = if dataset_bc !== nothing || + dataset_pde !== nothing + merge_strategy_with_loglikelihood_function(pinnrep, GridTraining(0.1), + datafree_pde_loss_functions, datafree_bc_loss_functions, + train_sets_pde = dataset_pde, train_sets_bc = dataset_bc) + else + nothing, nothing + end - # this includes losses from dataset domain points as well as discretization points - function full_loss_function(θ, allstd::Vector{Vector{Float64}}) - stdpdes, stdbcs, stdextra = allstd - # the aggregation happens on cpu even if the losses are gpu, probably fine since it's only a few of them - # SSE FOR LOSS ON GRIDPOINTS not MSE ! i, j depend on number of bcs and eqs - pde_loglikelihoods = sum([pde_loglike_function(θ, stdpdes[i]) - for (i, pde_loglike_function) in enumerate(pde_loss_functions)]) + # this includes losses from dataset domain points as well as discretization points + function full_loss_function(θ, allstd::Vector{Vector{Float64}}) + stdpdes, stdbcs, stdextra = allstd + # the aggregation happens on cpu even if the losses are gpu, probably fine since it's only a few of them + # SSE FOR LOSS ON GRIDPOINTS not MSE ! i, j depend on number of bcs and eqs + pde_loglikelihoods = sum([pde_loglike_function(θ, stdpdes[i]) + for (i, pde_loglike_function) in enumerate(pde_loss_functions)]) - bc_loglikelihoods = sum([bc_loglike_function(θ, stdbcs[j]) - for (j, bc_loglike_function) in enumerate(bc_loss_functions)]) + bc_loglikelihoods = sum([bc_loglike_function(θ, stdbcs[j]) + for (j, bc_loglike_function) in enumerate(bc_loss_functions)]) - # final newloss creation components are similar to this - if !(datapde_loss_functions isa Nothing) - pde_loglikelihoods += sum([pde_loglike_function(θ, stdpdes[j]) - for (j, pde_loglike_function) in enumerate(datapde_loss_functions)]) - end + # final newloss creation components are similar to this + if !(datapde_loss_functions isa Nothing) + pde_loglikelihoods += sum([pde_loglike_function(θ, stdpdes[j]) + for (j, pde_loglike_function) in enumerate(datapde_loss_functions)]) + end - if !(databc_loss_functions isa Nothing) - bc_loglikelihoods += sum([bc_loglike_function(θ, stdbcs[j]) - for (j, bc_loglike_function) in enumerate(databc_loss_functions)]) - end + if !(databc_loss_functions isa Nothing) + bc_loglikelihoods += sum([bc_loglike_function(θ, stdbcs[j]) + for (j, bc_loglike_function) in enumerate(databc_loss_functions)]) + end - # this is kind of a hack, and means that whenever the outer function is evaluated the increment goes up, even if it's not being optimized - # that's why we prefer the user to maintain the increment in the outer loop callback during optimization - @ignore_derivatives if self_increment - iteration[] += 1 - end + # this is kind of a hack, and means that whenever the outer function is evaluated the increment goes up, even if it's not being optimized + # that's why we prefer the user to maintain the increment in the outer loop callback during optimization + @ignore_derivatives if self_increment + iteration[] += 1 + end - @ignore_derivatives begin - reweight_losses_func(θ, pde_loglikelihoods, - bc_loglikelihoods) - end + @ignore_derivatives begin + reweight_losses_func(θ, pde_loglikelihoods, + bc_loglikelihoods) + end - weighted_pde_loglikelihood = adaloss.pde_loss_weights .* pde_loglikelihoods - weighted_bc_loglikelihood = adaloss.bc_loss_weights .* bc_loglikelihoods + weighted_pde_loglikelihood = adaloss.pde_loss_weights .* pde_loglikelihoods + weighted_bc_loglikelihood = adaloss.bc_loss_weights .* bc_loglikelihoods - sum_weighted_pde_loglikelihood = sum(weighted_pde_loglikelihood) - sum_weighted_bc_loglikelihood = sum(weighted_bc_loglikelihood) - weighted_loglikelihood_before_additional = sum_weighted_pde_loglikelihood + - sum_weighted_bc_loglikelihood + sum_weighted_pde_loglikelihood = sum(weighted_pde_loglikelihood) + sum_weighted_bc_loglikelihood = sum(weighted_bc_loglikelihood) + weighted_loglikelihood_before_additional = sum_weighted_pde_loglikelihood + + sum_weighted_bc_loglikelihood - full_weighted_loglikelihood = if additional_loss isa Nothing - weighted_loglikelihood_before_additional - else - (θ_, p_) = param_estim ? (θ.depvar, θ.p) : (θ, nothing) - _additional_loss = additional_loss(phi, θ_, p_) - _additional_loglikelihood = logpdf(Normal(0, stdextra), _additional_loss) + full_weighted_loglikelihood = if additional_loss isa Nothing + weighted_loglikelihood_before_additional + else + (θ_, p_) = param_estim ? (θ.depvar, θ.p) : (θ, nothing) + _additional_loss = additional_loss(phi, θ_, p_) + _additional_loglikelihood = logpdf(Normal(0, stdextra), _additional_loss) - weighted_additional_loglikelihood = adaloss.additional_loss_weights[1] * - _additional_loglikelihood + weighted_additional_loglikelihood = adaloss.additional_loss_weights[1] * + _additional_loglikelihood - weighted_loglikelihood_before_additional + weighted_additional_loglikelihood - end + weighted_loglikelihood_before_additional + weighted_additional_loglikelihood + end - return full_weighted_loglikelihood - end + return full_weighted_loglikelihood + end - return full_loss_function -end \ No newline at end of file + return full_loss_function +end diff --git a/lib/BayesianNeuralPDE/src/pinn_types.jl b/lib/BayesianNeuralPDE/src/pinn_types.jl index 65b365de5..c71194fd7 100644 --- a/lib/BayesianNeuralPDE/src/pinn_types.jl +++ b/lib/BayesianNeuralPDE/src/pinn_types.jl @@ -1,5 +1,5 @@ """ - BayesianPINN(args...; dataset = nothing, kwargs...) + 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 @@ -17,17 +17,17 @@ the ones mentioned below. inverse problem solving. """ @concrete struct BayesianPINN <: AbstractPINN - pinn <: PhysicsInformedNN - dataset + pinn <: PhysicsInformedNN + dataset::Any end function Base.getproperty(pinn::BayesianPINN, name::Symbol) - name === :dataset && return getfield(pinn, :dataset) - name === :pinn && return getfield(pinn, :pinn) - return getproperty(pinn.pinn, name) + name === :dataset && return getfield(pinn, :dataset) + name === :pinn && return getfield(pinn, :pinn) + return getproperty(pinn.pinn, name) end function BayesianPINN(args...; dataset = nothing, kwargs...) - dataset === nothing && (dataset = (nothing, nothing)) - return BayesianPINN(PhysicsInformedNN(args...; kwargs...), dataset) -end \ No newline at end of file + dataset === nothing && (dataset = (nothing, nothing)) + return BayesianPINN(PhysicsInformedNN(args...; kwargs...), dataset) +end diff --git a/lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl b/lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl index f5f5a96f7..296e4f216 100644 --- a/lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl +++ b/lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl @@ -1,529 +1,529 @@ -@testitem "BPINN PDE I: 1D Periodic System" tags=[:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum - - Random.seed!(100) - - @parameters t - @variables u(..) - Dt = Differential(t) - eq = Dt(u(t)) - cospi(2t) ~ 0 - bcs = [u(0.0) ~ 0.0] - domains = [t ∈ Interval(0.0, 2.0)] - - chainl = Chain(Dense(1, 6, tanh), Dense(6, 1)) - initl, st = Lux.setup(Random.default_rng(), chainl) - @named pde_system = PDESystem(eq, bcs, domains, [t], [u(t)]) - - # non adaptive case - discretization = BayesianPINN([chainl], GridTraining([0.01])) - - sol1 = ahmc_bayesian_pinn_pde( - pde_system, discretization; draw_samples = 1500, bcstd = [0.01], - phystd = [0.01], priorsNNw = (0.0, 1.0), saveats = [1 / 50.0]) - - analytic_sol_func(u0, t) = u0 + sinpi(2t) / (2pi) - ts = vec(sol1.timepoints[1]) - u_real = [analytic_sol_func(0.0, t) for t in ts] - u_predict = pmean(sol1.ensemblesol[1]) - - # absol tests - @test mean(abs, u_predict .- u_real) < 5e-2 +@testitem "BPINN PDE I: 1D Periodic System" tags = [:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum + + Random.seed!(100) + + @parameters t + @variables u(..) + Dt = Differential(t) + eq = Dt(u(t)) - cospi(2t) ~ 0 + bcs = [u(0.0) ~ 0.0] + domains = [t ∈ Interval(0.0, 2.0)] + + chainl = Chain(Dense(1, 6, tanh), Dense(6, 1)) + initl, st = Lux.setup(Random.default_rng(), chainl) + @named pde_system = PDESystem(eq, bcs, domains, [t], [u(t)]) + + # non adaptive case + discretization = BayesianPINN([chainl], GridTraining([0.01])) + + sol1 = ahmc_bayesian_pinn_pde( + pde_system, discretization; draw_samples = 1500, bcstd = [0.01], + phystd = [0.01], priorsNNw = (0.0, 1.0), saveats = [1 / 50.0]) + + analytic_sol_func(u0, t) = u0 + sinpi(2t) / (2pi) + ts = vec(sol1.timepoints[1]) + u_real = [analytic_sol_func(0.0, t) for t in ts] + u_predict = pmean(sol1.ensemblesol[1]) + + # absol tests + @test mean(abs, u_predict .- u_real) < 5e-2 end -@testitem "BPINN PDE II: 1D ODE" tags=[:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum +@testitem "BPINN PDE II: 1D ODE" tags = [:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum - Random.seed!(100) + Random.seed!(100) - @parameters θ - @variables u(..) - Dθ = Differential(θ) + @parameters θ + @variables u(..) + Dθ = Differential(θ) - # 1D ODE - eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) - - u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3))) + # 1D ODE + eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) - + u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3))) - # Initial and boundary conditions - bcs = [u(0.0) ~ 1.0f0] + # Initial and boundary conditions + bcs = [u(0.0) ~ 1.0f0] - # Space and time domains - domains = [θ ∈ Interval(0.0f0, 1.0f0)] + # Space and time domains + domains = [θ ∈ Interval(0.0f0, 1.0f0)] - # Discretization - dt = 0.1f0 + # Discretization + dt = 0.1f0 - # Neural network - chain = Chain(Dense(1, 12, σ), Dense(12, 1)) + # Neural network + chain = Chain(Dense(1, 12, σ), Dense(12, 1)) - discretization = BayesianPINN([chain], GridTraining([0.01])) + discretization = BayesianPINN([chain], GridTraining([0.01])) - @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) + @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) - sol1 = ahmc_bayesian_pinn_pde( - pde_system, discretization; draw_samples = 500, bcstd = [0.1], - phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) + sol1 = ahmc_bayesian_pinn_pde( + pde_system, discretization; draw_samples = 500, bcstd = [0.1], + phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) - analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2 - ts = sol1.timepoints[1] - u_real = vec([analytic_sol_func(t) for t in ts]) - u_predict = pmean(sol1.ensemblesol[1]) - @test u_predict≈u_real atol=0.8 + analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2 + ts = sol1.timepoints[1] + u_real = vec([analytic_sol_func(t) for t in ts]) + u_predict = pmean(sol1.ensemblesol[1]) + @test u_predict ≈ u_real atol = 0.8 end -@testitem "BPINN PDE III: 3rd Degree ODE" tags=[:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum +@testitem "BPINN PDE III: 3rd Degree ODE" tags = [:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum - Random.seed!(100) + Random.seed!(100) - @parameters x - @variables u(..), Dxu(..), Dxxu(..), O1(..), O2(..) - Dxxx = Differential(x)^3 - Dx = Differential(x) + @parameters x + @variables u(..), Dxu(..), Dxxu(..), O1(..), O2(..) + Dxxx = Differential(x)^3 + Dx = Differential(x) - # ODE - eq = Dx(Dxxu(x)) ~ cospi(x) + # ODE + eq = Dx(Dxxu(x)) ~ cospi(x) - # Initial and boundary conditions - ep = (cbrt(eps(eltype(Float64))))^2 / 6 + # Initial and boundary conditions + ep = (cbrt(eps(eltype(Float64))))^2 / 6 - bcs = [ - u(0.0) ~ 0.0, - u(1.0) ~ cospi(1.0), - Dxu(1.0) ~ 1.0, - Dxu(x) ~ Dx(u(x)) + ep * O1(x), - Dxxu(x) ~ Dx(Dxu(x)) + ep * O2(x) - ] + bcs = [ + u(0.0) ~ 0.0, + u(1.0) ~ cospi(1.0), + Dxu(1.0) ~ 1.0, + Dxu(x) ~ Dx(u(x)) + ep * O1(x), + Dxxu(x) ~ Dx(Dxu(x)) + ep * O2(x), + ] - # Space and time domains - domains = [x ∈ Interval(0.0, 1.0)] + # Space and time domains + domains = [x ∈ Interval(0.0, 1.0)] - # Neural network - chain = [ - Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), - Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), - Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), - Chain(Dense(1, 4, tanh), Dense(4, 1)), - Chain(Dense(1, 4, tanh), Dense(4, 1)) - ] + # Neural network + chain = [ + Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), + Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), + Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), + Chain(Dense(1, 4, tanh), Dense(4, 1)), + Chain(Dense(1, 4, tanh), Dense(4, 1)), + ] - discretization = BayesianPINN(chain, GridTraining(0.01)) + discretization = BayesianPINN(chain, GridTraining(0.01)) - @named pde_system = PDESystem(eq, bcs, domains, [x], - [u(x), Dxu(x), Dxxu(x), O1(x), O2(x)]) + @named pde_system = PDESystem(eq, bcs, domains, [x], + [u(x), Dxu(x), Dxxu(x), O1(x), O2(x)]) - sol1 = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200, - bcstd = [0.01, 0.01, 0.01, 0.01, 0.01], phystd = [0.005], - priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) + sol1 = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200, + bcstd = [0.01, 0.01, 0.01, 0.01, 0.01], phystd = [0.005], + priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) - analytic_sol_func(x) = (π * x * (-x + (π^2) * (2 * x - 3) + 1) - sinpi(x)) / (π^3) + analytic_sol_func(x) = (π * x * (-x + (π^2) * (2 * x - 3) + 1) - sinpi(x)) / (π^3) - u_predict = pmean(sol1.ensemblesol[1]) - xs = vec(sol1.timepoints[1]) - u_real = [analytic_sol_func(x) for x in xs] - @test u_predict≈u_real atol=0.5 + u_predict = pmean(sol1.ensemblesol[1]) + xs = vec(sol1.timepoints[1]) + u_real = [analytic_sol_func(x) for x in xs] + @test u_predict ≈ u_real atol = 0.5 end -@testitem "BPINN PDE IV: 2D Poisson" tags=[:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum +@testitem "BPINN PDE IV: 2D Poisson" tags = [:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum - Random.seed!(100) + Random.seed!(100) - @parameters x y - @variables u(..) - Dxx = Differential(x)^2 - Dyy = Differential(y)^2 + @parameters x y + @variables u(..) + Dxx = Differential(x)^2 + Dyy = Differential(y)^2 - # 2D PDE - eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y) + # 2D PDE + eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y) - # Boundary conditions - bcs = [ - u(0, y) ~ 0.0, - u(1, y) ~ 0.0, - u(x, 0) ~ 0.0, - u(x, 1) ~ 0.0 - ] + # Boundary conditions + bcs = [ + u(0, y) ~ 0.0, + u(1, y) ~ 0.0, + u(x, 0) ~ 0.0, + u(x, 1) ~ 0.0, + ] - # Space and time domains - domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)] + # Space and time domains + domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)] - # Discretization - dt = 0.1f0 + # Discretization + dt = 0.1f0 - # Neural network - chain = Chain(Dense(2, 9, σ), Dense(9, 9, σ), Dense(9, 1)) + # Neural network + chain = Chain(Dense(2, 9, σ), Dense(9, 9, σ), Dense(9, 1)) - dx = 0.04 - discretization = BayesianPINN([chain], GridTraining(dx)) + dx = 0.04 + discretization = BayesianPINN([chain], GridTraining(dx)) - @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) + @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) - sol = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200, - bcstd = [0.003, 0.003, 0.003, 0.003], phystd = [0.003], - priorsNNw = (0.0, 10.0), saveats = [1 / 100.0, 1 / 100.0]) + sol = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200, + bcstd = [0.003, 0.003, 0.003, 0.003], phystd = [0.003], + priorsNNw = (0.0, 10.0), saveats = [1 / 100.0, 1 / 100.0]) - xs = sol.timepoints[1] - analytic_sol_func(x, y) = (sinpi(x) * sinpi(y)) / (2pi^2) + xs = sol.timepoints[1] + analytic_sol_func(x, y) = (sinpi(x) * sinpi(y)) / (2pi^2) - u_predict = pmean(sol.ensemblesol[1]) - u_real = [analytic_sol_func(xs[:, i][1], xs[:, i][2]) for i in 1:length(xs[1, :])] - @test u_predict≈u_real rtol=0.5 + u_predict = pmean(sol.ensemblesol[1]) + u_real = [analytic_sol_func(xs[:, i][1], xs[:, i][2]) for i in 1:length(xs[1, :])] + @test u_predict ≈ u_real rtol = 0.5 end -@testitem "BPINN PDE: Translating from Flux" tags=[:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum - import Flux +@testitem "BPINN PDE: Translating from Flux" tags = [:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum + using Flux: Flux - Random.seed!(100) + Random.seed!(100) - @parameters θ - @variables u(..) - Dθ = Differential(θ) + @parameters θ + @variables u(..) + Dθ = Differential(θ) - # 1D ODE - eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) - - u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3))) + # 1D ODE + eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) - + u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3))) - # Initial and boundary conditions - bcs = [u(0.0) ~ 1.0f0] + # Initial and boundary conditions + bcs = [u(0.0) ~ 1.0f0] - # Space and time domains - domains = [θ ∈ Interval(0.0f0, 1.0f0)] + # Space and time domains + domains = [θ ∈ Interval(0.0f0, 1.0f0)] - # Neural network - chain = Flux.Chain(Flux.Dense(1, 12, Flux.σ), Flux.Dense(12, 1)) + # Neural network + chain = Flux.Chain(Flux.Dense(1, 12, Flux.σ), Flux.Dense(12, 1)) - discretization = BayesianPINN([chain], GridTraining([0.01])) - @test discretization.chain[1] isa Lux.AbstractLuxLayer + discretization = BayesianPINN([chain], GridTraining([0.01])) + @test discretization.chain[1] isa Lux.AbstractLuxLayer - @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) + @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) - sol = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 500, - bcstd = [0.1], phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) + sol = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 500, + bcstd = [0.1], phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) - analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2 - ts = sol.timepoints[1] - u_real = vec([analytic_sol_func(t) for t in ts]) - u_predict = pmean(sol.ensemblesol[1]) + analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2 + ts = sol.timepoints[1] + u_real = vec([analytic_sol_func(t) for t in ts]) + u_predict = pmean(sol.ensemblesol[1]) - @test u_predict≈u_real atol=0.8 + @test u_predict ≈ u_real atol = 0.8 end -@testitem "BPINN PDE Inv I: 1D Periodic System" tags=[:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum - - Random.seed!(100) - - @parameters t p - @variables u(..) - - Dt = Differential(t) - eqs = Dt(u(t)) - cos(p * t) ~ 0 - bcs = [u(0) ~ 0.0] - domains = [t ∈ Interval(0.0, 2.0)] - - chainl = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) - initl, st = Lux.setup(Random.default_rng(), chainl) - - @named pde_system = PDESystem(eqs, - bcs, - domains, - [t], - [u(t)], - [p], - defaults = Dict([p => 4.0])) - - analytic_sol_func1(u0, t) = u0 + sinpi(2t) / (2π) - timepoints = collect(0.0:(1 / 100.0):2.0) - u = [analytic_sol_func1(0.0, timepoint) for timepoint in timepoints] - u = u .+ (u .* 0.2) .* randn(size(u)) - dataset = [hcat(u, timepoints)] - - # BPINNs are formulated with a mesh that must stay the same throughout sampling (as of now) - @testset "$(nameof(typeof(strategy)))" for strategy in [ - # StochasticTraining(200), - # QuasiRandomTraining(200), - GridTraining([0.02]) - ] - discretization = BayesianPINN([chainl], strategy; param_estim = true, - dataset = [dataset, nothing]) - - sol1 = ahmc_bayesian_pinn_pde(pde_system, - discretization; - draw_samples = 1500, - bcstd = [0.02], - phystd = [0.02], l2std = [0.02], - priorsNNw = (0.0, 1.0), - saveats = [1 / 50.0], - param = [LogNormal(6.0, 0.5)]) - - param = 2 * π - ts = vec(sol1.timepoints[1]) - u_real = [analytic_sol_func1(0.0, t) for t in ts] - u_predict = pmean(sol1.ensemblesol[1]) - - @test mean(abs, u_predict .- u_real) < 5e-2 - @test sol1.estimated_de_params[1]≈param rtol=0.1 - end +@testitem "BPINN PDE Inv I: 1D Periodic System" tags = [:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum + + Random.seed!(100) + + @parameters t p + @variables u(..) + + Dt = Differential(t) + eqs = Dt(u(t)) - cos(p * t) ~ 0 + bcs = [u(0) ~ 0.0] + domains = [t ∈ Interval(0.0, 2.0)] + + chainl = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) + initl, st = Lux.setup(Random.default_rng(), chainl) + + @named pde_system = PDESystem(eqs, + bcs, + domains, + [t], + [u(t)], + [p], + defaults = Dict([p => 4.0])) + + analytic_sol_func1(u0, t) = u0 + sinpi(2t) / (2π) + timepoints = collect(0.0:(1/100.0):2.0) + u = [analytic_sol_func1(0.0, timepoint) for timepoint in timepoints] + u = u .+ (u .* 0.2) .* randn(size(u)) + dataset = [hcat(u, timepoints)] + + # BPINNs are formulated with a mesh that must stay the same throughout sampling (as of now) + @testset "$(nameof(typeof(strategy)))" for strategy in [ + # StochasticTraining(200), + # QuasiRandomTraining(200), + GridTraining([0.02]), + ] + discretization = BayesianPINN([chainl], strategy; param_estim = true, + dataset = [dataset, nothing]) + + sol1 = ahmc_bayesian_pinn_pde(pde_system, + discretization; + draw_samples = 1500, + bcstd = [0.02], + phystd = [0.02], l2std = [0.02], + priorsNNw = (0.0, 1.0), + saveats = [1 / 50.0], + param = [LogNormal(6.0, 0.5)]) + + param = 2 * π + ts = vec(sol1.timepoints[1]) + u_real = [analytic_sol_func1(0.0, t) for t in ts] + u_predict = pmean(sol1.ensemblesol[1]) + + @test mean(abs, u_predict .- u_real) < 5e-2 + @test sol1.estimated_de_params[1] ≈ param rtol = 0.1 + end end -@testitem "BPINN PDE Inv II: Lorenz System" tags=[:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum - - Random.seed!(100) - - @parameters t, σ_ - @variables x(..), y(..), z(..) - Dt = Differential(t) - eqs = [ - Dt(x(t)) ~ σ_ * (y(t) - x(t)), - Dt(y(t)) ~ x(t) * (28.0 - z(t)) - y(t), - Dt(z(t)) ~ x(t) * y(t) - 8.0 / 3.0 * z(t) - ] - - bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0] - domains = [t ∈ Interval(0.0, 1.0)] - - input_ = length(domains) - n = 7 - chain = [ - Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)), - Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)), - Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)) - ] - - # Generate Data - function lorenz!(du, u, p, t) - du[1] = 10.0 * (u[2] - u[1]) - du[2] = u[1] * (28.0 - u[3]) - u[2] - du[3] = u[1] * u[2] - (8.0 / 3.0) * u[3] - end - - u0 = [1.0; 0.0; 0.0] - tspan = (0.0, 1.0) - prob = ODEProblem(lorenz!, u0, tspan) - sol = solve(prob, Tsit5(), dt = 0.01, saveat = 0.05) - ts = sol.t - us = hcat(sol.u...) - us = us .+ ((0.05 .* randn(size(us))) .* us) - ts_ = hcat(ts...)[1, :] - dataset = [hcat(us[i, :], ts_) for i in 1:3] - - discretization = BayesianPINN(chain, GridTraining([0.01]); param_estim = true, - dataset = [dataset, nothing]) - - @named pde_system = PDESystem(eqs, bcs, domains, - [t], [x(t), y(t), z(t)], [σ_], defaults = Dict([p => 1.0 for p in [σ_]])) - - sol1 = ahmc_bayesian_pinn_pde(pde_system, - discretization; - draw_samples = 50, - bcstd = [0.3, 0.3, 0.3], - phystd = [0.1, 0.1, 0.1], - l2std = [1, 1, 1], - priorsNNw = (0.0, 1.0), - saveats = [0.01], - param = [Normal(12.0, 2)]) - - idealp = 10.0 - p_ = sol1.estimated_de_params[1] - @test sum(abs, pmean(p_) - 10.00) < 0.3 * idealp[1] - # @test sum(abs, pmean(p_[2]) - (8 / 3)) < 0.3 * idealp[2] +@testitem "BPINN PDE Inv II: Lorenz System" tags = [:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum + + Random.seed!(100) + + @parameters t, σ_ + @variables x(..), y(..), z(..) + Dt = Differential(t) + eqs = [ + Dt(x(t)) ~ σ_ * (y(t) - x(t)), + Dt(y(t)) ~ x(t) * (28.0 - z(t)) - y(t), + Dt(z(t)) ~ x(t) * y(t) - 8.0 / 3.0 * z(t), + ] + + bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0] + domains = [t ∈ Interval(0.0, 1.0)] + + input_ = length(domains) + n = 7 + chain = [ + Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)), + Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)), + Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)), + ] + + # Generate Data + function lorenz!(du, u, p, t) + du[1] = 10.0 * (u[2] - u[1]) + du[2] = u[1] * (28.0 - u[3]) - u[2] + du[3] = u[1] * u[2] - (8.0 / 3.0) * u[3] + end + + u0 = [1.0; 0.0; 0.0] + tspan = (0.0, 1.0) + prob = ODEProblem(lorenz!, u0, tspan) + sol = solve(prob, Tsit5(), dt = 0.01, saveat = 0.05) + ts = sol.t + us = hcat(sol.u...) + us = us .+ ((0.05 .* randn(size(us))) .* us) + ts_ = hcat(ts...)[1, :] + dataset = [hcat(us[i, :], ts_) for i in 1:3] + + discretization = BayesianPINN(chain, GridTraining([0.01]); param_estim = true, + dataset = [dataset, nothing]) + + @named pde_system = PDESystem(eqs, bcs, domains, + [t], [x(t), y(t), z(t)], [σ_], defaults = Dict([p => 1.0 for p in [σ_]])) + + sol1 = ahmc_bayesian_pinn_pde(pde_system, + discretization; + draw_samples = 50, + bcstd = [0.3, 0.3, 0.3], + phystd = [0.1, 0.1, 0.1], + l2std = [1, 1, 1], + priorsNNw = (0.0, 1.0), + saveats = [0.01], + param = [Normal(12.0, 2)]) + + idealp = 10.0 + p_ = sol1.estimated_de_params[1] + @test sum(abs, pmean(p_) - 10.00) < 0.3 * idealp[1] + # @test sum(abs, pmean(p_[2]) - (8 / 3)) < 0.3 * idealp[2] end -@testitem "BPINN PDE Inv III: Improved Parametric Kuromo-Sivashinsky Equation solve" tags=[:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum - - Random.seed!(100) - - function recur_expression(exp, Dict_differentials) - for in_exp in exp.args - if !(in_exp isa Expr) - # skip +,== symbols, characters etc - continue - - elseif in_exp.args[1] isa ModelingToolkit.Differential - # first symbol of differential term - # Dict_differentials for masking differential terms - # and resubstituting differentials in equations after putting in interpolations - # temp = in_exp.args[end] - Dict_differentials[eval(in_exp)] = Symbolics.variable("diff_$(length(Dict_differentials) + 1)") - return - else - recur_expression(in_exp, Dict_differentials) - end - end - end - - @parameters x, t, α - @variables u(..) - Dt = Differential(t) - Dx = Differential(x) - Dx2 = Differential(x)^2 - Dx3 = Differential(x)^3 - Dx4 = Differential(x)^4 - - # α = 1 (KS equation to be parametric in a) - β = 4 - γ = 1 - eq = Dt(u(x, t)) + u(x, t) * Dx(u(x, t)) + α * Dx2(u(x, t)) + β * Dx3(u(x, t)) + γ * Dx4(u(x, t)) ~ 0 - - u_analytic(x, t; z = -x / 2 + t) = 11 + 15 * tanh(z) - 15 * tanh(z)^2 - 15 * tanh(z)^3 - du(x, t; z = -x / 2 + t) = 15 / 2 * (tanh(z) + 1) * (3 * tanh(z) - 1) * sech(z)^2 - - bcs = [u(x, 0) ~ u_analytic(x, 0), - u(-10, t) ~ u_analytic(-10, t), - u(10, t) ~ u_analytic(10, t), - Dx(u(-10, t)) ~ du(-10, t), - Dx(u(10, t)) ~ du(10, t)] - - # Space and time domains - domains = [x ∈ Interval(-10.0, 10.0), - t ∈ Interval(0.0, 1.0)] - - # Discretization - dx = 0.4 - dt = 0.2 - - # Function to compute analytical solution at a specific point (x, t) - function u_analytic_point(x, t) - z = -x / 2 + t - return 11 + 15 * tanh(z) - 15 * tanh(z)^2 - 15 * tanh(z)^3 - end - - # Function to generate the dataset matrix - function generate_dataset_matrix(domains, dx, dt, xlim, tlim) - x_values = xlim[1]:dx:xlim[2] - t_values = tlim[1]:dt:tlim[2] - - dataset = [] - - for t in t_values - for x in x_values - u_value = u_analytic_point(x, t) - push!(dataset, [u_value, x, t]) - end - end - - return vcat([data' for data in dataset]...) - end - - # considering sparse dataset from half of x's domain - datasetpde_new = [generate_dataset_matrix(domains, dx, dt, [-10, 0], [0.0, 1.0])] - - # Adding Gaussian noise with a 0.8 std - noisydataset_new = deepcopy(datasetpde_new) - noisydataset_new[1][:, 1] = noisydataset_new[1][:, 1] .+ - (randn(size(noisydataset_new[1][:, 1])) .* 0.8) - - # Neural network - chain = Lux.Chain(Lux.Dense(2, 8, Lux.tanh), - Lux.Dense(8, 8, Lux.tanh), - Lux.Dense(8, 1)) - - # Discretization for old and new models - discretization = NeuralPDE.BayesianPINN([chain], - GridTraining([dx, dt]), param_estim = true, dataset = [noisydataset_new, nothing]) - - # let α default to 2.0 - @named pde_system = PDESystem(eq, - bcs, - domains, - [x, t], - [u(x, t)], - [α], - defaults = Dict([α => 2.0])) - - # neccesarry for loss function construction (involves Operator masking) - eqs = pde_system.eqs - Dict_differentials = Dict() - exps = toexpr.(eqs) - nullobj = [recur_expression(exp, Dict_differentials) for exp in exps] - - # Dict_differentials is now ; - # Dict{Any, Any} with 5 entries: - # Differential(x)(Differential(x)(u(x, t))) => diff_5 - # Differential(x)(Differential(x)(Differential(x)(u(x… => diff_1 - # Differential(x)(Differential(x)(Differential(x)(Dif… => diff_2 - # Differential(x)(u(x, t)) => diff_4 - # Differential(t)(u(x, t)) => diff_3 - - # using HMC algorithm due to convergence, stability, time of training. (refer to mcmc chain plots) - # choice of std for objectives is very important - # pass in Dict_differentials, phystdnew arguments when using the new model - - sol_new = ahmc_bayesian_pinn_pde(pde_system, - discretization; - draw_samples = 150, - bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], phynewstd = [0.4], - phystd = [0.2], l2std = [0.8], param = [Distributions.Normal(2.0, 2)], - priorsNNw = (0.0, 1.0), - saveats = [1 / 100.0, 1 / 100.0], - Dict_differentials = Dict_differentials) - - sol_old = ahmc_bayesian_pinn_pde(pde_system, - discretization; - draw_samples = 150, - bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], - phystd = [0.2], l2std = [0.8], param = [Distributions.Normal(2.0, 2)], - priorsNNw = (0.0, 1.0), - saveats = [1 / 100.0, 1 / 100.0]) - - phi = discretization.phi[1] - xs, ts = [infimum(d.domain):dx:supremum(d.domain) - for (d, dx) in zip(domains, [dx / 10, dt])] - u_real = [[u_analytic(x, t) for x in xs] for t in ts] - - u_predict_new = [[first(pmean(phi([x, t], sol_new.estimated_nn_params[1]))) for x in xs] - for t in ts] - - diff_u_new = [[abs(u_analytic(x, t) - - first(pmean(phi([x, t], sol_new.estimated_nn_params[1])))) - for x in xs] - for t in ts] - - u_predict_old = [[first(pmean(phi([x, t], sol_old.estimated_nn_params[1]))) for x in xs] - for t in ts] - diff_u_old = [[abs(u_analytic(x, t) - - first(pmean(phi([x, t], sol_old.estimated_nn_params[1])))) - for x in xs] - for t in ts] - - unsafe_comparisons(true) - @test all(all, [((diff_u_new[i]) .^ 2 .< 0.8) for i in 1:6]) == true - @test all(all, [((diff_u_old[i]) .^ 2 .< 0.8) for i in 1:6]) == false - - MSE_new = [mean(abs2, diff_u_new[i]) for i in 1:6] - MSE_old = [mean(abs2, diff_u_old[i]) for i in 1:6] - @test (MSE_new .< MSE_old) == [1, 1, 1, 1, 1, 1] - - param_new = sol_new.estimated_de_params[1] - param_old = sol_old.estimated_de_params[1] - α = 1 - @test abs(param_new - α) < 0.2 * α - @test abs(param_new - α) < abs(param_old - α) -end \ No newline at end of file +@testitem "BPINN PDE Inv III: Improved Parametric Kuromo-Sivashinsky Equation solve" tags = [:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum + + Random.seed!(100) + + function recur_expression(exp, Dict_differentials) + for in_exp in exp.args + if !(in_exp isa Expr) + # skip +,== symbols, characters etc + continue + + elseif in_exp.args[1] isa ModelingToolkit.Differential + # first symbol of differential term + # Dict_differentials for masking differential terms + # and resubstituting differentials in equations after putting in interpolations + # temp = in_exp.args[end] + Dict_differentials[eval(in_exp)] = Symbolics.variable("diff_$(length(Dict_differentials) + 1)") + return + else + recur_expression(in_exp, Dict_differentials) + end + end + end + + @parameters x, t, α + @variables u(..) + Dt = Differential(t) + Dx = Differential(x) + Dx2 = Differential(x)^2 + Dx3 = Differential(x)^3 + Dx4 = Differential(x)^4 + + # α = 1 (KS equation to be parametric in a) + β = 4 + γ = 1 + eq = Dt(u(x, t)) + u(x, t) * Dx(u(x, t)) + α * Dx2(u(x, t)) + β * Dx3(u(x, t)) + γ * Dx4(u(x, t)) ~ 0 + + u_analytic(x, t; z = -x / 2 + t) = 11 + 15 * tanh(z) - 15 * tanh(z)^2 - 15 * tanh(z)^3 + du(x, t; z = -x / 2 + t) = 15 / 2 * (tanh(z) + 1) * (3 * tanh(z) - 1) * sech(z)^2 + + bcs = [u(x, 0) ~ u_analytic(x, 0), + u(-10, t) ~ u_analytic(-10, t), + u(10, t) ~ u_analytic(10, t), + Dx(u(-10, t)) ~ du(-10, t), + Dx(u(10, t)) ~ du(10, t)] + + # Space and time domains + domains = [x ∈ Interval(-10.0, 10.0), + t ∈ Interval(0.0, 1.0)] + + # Discretization + dx = 0.4 + dt = 0.2 + + # Function to compute analytical solution at a specific point (x, t) + function u_analytic_point(x, t) + z = -x / 2 + t + return 11 + 15 * tanh(z) - 15 * tanh(z)^2 - 15 * tanh(z)^3 + end + + # Function to generate the dataset matrix + function generate_dataset_matrix(domains, dx, dt, xlim, tlim) + x_values = xlim[1]:dx:xlim[2] + t_values = tlim[1]:dt:tlim[2] + + dataset = [] + + for t in t_values + for x in x_values + u_value = u_analytic_point(x, t) + push!(dataset, [u_value, x, t]) + end + end + + return vcat([data' for data in dataset]...) + end + + # considering sparse dataset from half of x's domain + datasetpde_new = [generate_dataset_matrix(domains, dx, dt, [-10, 0], [0.0, 1.0])] + + # Adding Gaussian noise with a 0.8 std + noisydataset_new = deepcopy(datasetpde_new) + noisydataset_new[1][:, 1] = noisydataset_new[1][:, 1] .+ + (randn(size(noisydataset_new[1][:, 1])) .* 0.8) + + # Neural network + chain = Lux.Chain(Lux.Dense(2, 8, Lux.tanh), + Lux.Dense(8, 8, Lux.tanh), + Lux.Dense(8, 1)) + + # Discretization for old and new models + discretization = NeuralPDE.BayesianPINN([chain], + GridTraining([dx, dt]), param_estim = true, dataset = [noisydataset_new, nothing]) + + # let α default to 2.0 + @named pde_system = PDESystem(eq, + bcs, + domains, + [x, t], + [u(x, t)], + [α], + defaults = Dict([α => 2.0])) + + # neccesarry for loss function construction (involves Operator masking) + eqs = pde_system.eqs + Dict_differentials = Dict() + exps = toexpr.(eqs) + nullobj = [recur_expression(exp, Dict_differentials) for exp in exps] + + # Dict_differentials is now ; + # Dict{Any, Any} with 5 entries: + # Differential(x)(Differential(x)(u(x, t))) => diff_5 + # Differential(x)(Differential(x)(Differential(x)(u(x… => diff_1 + # Differential(x)(Differential(x)(Differential(x)(Dif… => diff_2 + # Differential(x)(u(x, t)) => diff_4 + # Differential(t)(u(x, t)) => diff_3 + + # using HMC algorithm due to convergence, stability, time of training. (refer to mcmc chain plots) + # choice of std for objectives is very important + # pass in Dict_differentials, phystdnew arguments when using the new model + + sol_new = ahmc_bayesian_pinn_pde(pde_system, + discretization; + draw_samples = 150, + bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], phynewstd = [0.4], + phystd = [0.2], l2std = [0.8], param = [Distributions.Normal(2.0, 2)], + priorsNNw = (0.0, 1.0), + saveats = [1 / 100.0, 1 / 100.0], + Dict_differentials = Dict_differentials) + + sol_old = ahmc_bayesian_pinn_pde(pde_system, + discretization; + draw_samples = 150, + bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], + phystd = [0.2], l2std = [0.8], param = [Distributions.Normal(2.0, 2)], + priorsNNw = (0.0, 1.0), + saveats = [1 / 100.0, 1 / 100.0]) + + phi = discretization.phi[1] + xs, ts = [infimum(d.domain):dx:supremum(d.domain) + for (d, dx) in zip(domains, [dx / 10, dt])] + u_real = [[u_analytic(x, t) for x in xs] for t in ts] + + u_predict_new = [[first(pmean(phi([x, t], sol_new.estimated_nn_params[1]))) for x in xs] + for t in ts] + + diff_u_new = [[abs(u_analytic(x, t) - + first(pmean(phi([x, t], sol_new.estimated_nn_params[1])))) + for x in xs] + for t in ts] + + u_predict_old = [[first(pmean(phi([x, t], sol_old.estimated_nn_params[1]))) for x in xs] + for t in ts] + diff_u_old = [[abs(u_analytic(x, t) - + first(pmean(phi([x, t], sol_old.estimated_nn_params[1])))) + for x in xs] + for t in ts] + + unsafe_comparisons(true) + @test all(all, [((diff_u_new[i]) .^ 2 .< 0.8) for i in 1:6]) == true + @test all(all, [((diff_u_old[i]) .^ 2 .< 0.8) for i in 1:6]) == false + + MSE_new = [mean(abs2, diff_u_new[i]) for i in 1:6] + MSE_old = [mean(abs2, diff_u_old[i]) for i in 1:6] + @test (MSE_new .< MSE_old) == [1, 1, 1, 1, 1, 1] + + param_new = sol_new.estimated_de_params[1] + param_old = sol_old.estimated_de_params[1] + α = 1 + @test abs(param_new - α) < 0.2 * α + @test abs(param_new - α) < abs(param_old - α) +end diff --git a/lib/BayesianNeuralPDE/test/BPINN_tests.jl b/lib/BayesianNeuralPDE/test/BPINN_tests.jl index f2a33355a..494228275 100644 --- a/lib/BayesianNeuralPDE/test/BPINN_tests.jl +++ b/lib/BayesianNeuralPDE/test/BPINN_tests.jl @@ -1,429 +1,429 @@ -@testitem "BPINN ODE I: Without Param Estimation" tags=[:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - import Flux - - Random.seed!(100) - - linear_analytic = (u0, p, t) -> u0 + sin(2 * π * t) / (2 * π) - linear = (u, p, t) -> cos(2 * π * t) - tspan = (0.0, 2.0) - u0 = 0.0 - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan) - p = prob.p - - # Numerical and Analytical Solutions: testing ahmc_bayesian_pinn_ode() - ta = range(tspan[1], tspan[2], length = 300) - u = [linear_analytic(u0, nothing, ti) for ti in ta] - x̂ = collect(Float64, Array(u) + 0.02 * randn(size(u))) - time = vec(collect(Float64, ta)) - physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] - - # testing points for solve() call must match saveat(1/50.0) arg - ta0 = range(tspan[1], tspan[2], length = 101) - u1 = [linear_analytic(u0, nothing, ti) for ti in ta0] - x̂1 = collect(Float64, Array(u1) + 0.02 * randn(size(u1))) - time1 = vec(collect(Float64, ta0)) - physsol0_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] - - chainlux = Chain(Dense(1, 7, tanh), Dense(7, 1)) - θinit, st = Lux.setup(Random.default_rng(), chainlux) - - fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( - prob, chainlux, draw_samples = 2500) - - alg = BNNODE(chainlux, draw_samples = 2500) - sol1lux = solve(prob, alg) - - # testing points - t = time - # Mean of last 500 sampled parameter's curves[Ensemble predictions] - θ = [vector_to_parameters(fhsamples[i], θinit) for i in 2000:length(fhsamples)] - luxar = [chainlux(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - # --------------------- ahmc_bayesian_pinn_ode() call - @test mean(abs.(x̂ .- meanscurve)) < 0.05 - @test mean(abs.(physsol1 .- meanscurve)) < 0.005 - - #--------------------- solve() call - @test mean(abs.(x̂1 .- pmean(sol1lux.ensemblesol[1]))) < 0.025 - @test mean(abs.(physsol0_1 .- pmean(sol1lux.ensemblesol[1]))) < 0.025 +@testitem "BPINN ODE I: Without Param Estimation" tags = [:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear_analytic = (u0, p, t) -> u0 + sin(2 * π * t) / (2 * π) + linear = (u, p, t) -> cos(2 * π * t) + tspan = (0.0, 2.0) + u0 = 0.0 + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan) + p = prob.p + + # Numerical and Analytical Solutions: testing ahmc_bayesian_pinn_ode() + ta = range(tspan[1], tspan[2], length = 300) + u = [linear_analytic(u0, nothing, ti) for ti in ta] + x̂ = collect(Float64, Array(u) + 0.02 * randn(size(u))) + time = vec(collect(Float64, ta)) + physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] + + # testing points for solve() call must match saveat(1/50.0) arg + ta0 = range(tspan[1], tspan[2], length = 101) + u1 = [linear_analytic(u0, nothing, ti) for ti in ta0] + x̂1 = collect(Float64, Array(u1) + 0.02 * randn(size(u1))) + time1 = vec(collect(Float64, ta0)) + physsol0_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] + + chainlux = Chain(Dense(1, 7, tanh), Dense(7, 1)) + θinit, st = Lux.setup(Random.default_rng(), chainlux) + + fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( + prob, chainlux, draw_samples = 2500) + + alg = BNNODE(chainlux, draw_samples = 2500) + sol1lux = solve(prob, alg) + + # testing points + t = time + # Mean of last 500 sampled parameter's curves[Ensemble predictions] + θ = [vector_to_parameters(fhsamples[i], θinit) for i in 2000:length(fhsamples)] + luxar = [chainlux(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + # --------------------- ahmc_bayesian_pinn_ode() call + @test mean(abs.(x̂ .- meanscurve)) < 0.05 + @test mean(abs.(physsol1 .- meanscurve)) < 0.005 + + #--------------------- solve() call + @test mean(abs.(x̂1 .- pmean(sol1lux.ensemblesol[1]))) < 0.025 + @test mean(abs.(physsol0_1 .- pmean(sol1lux.ensemblesol[1]))) < 0.025 end -@testitem "BPINN ODE II: With Parameter Estimation" tags=[:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - import Flux - - Random.seed!(100) - - linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - linear = (u, p, t) -> cos(p * t) - tspan = (0.0, 2.0) - u0 = 0.0 - p = 2 * pi - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, p) - - # Numerical and Analytical Solutions - sol1 = solve(prob, Tsit5(); saveat = 0.01) - u = sol1.u - time = sol1.t - - # BPINN AND TRAINING DATASET CREATION(dataset must be defined only inside problem timespan!) - ta = range(tspan[1], tspan[2], length = 100) - u = [linear_analytic(u0, p, ti) for ti in ta] - x̂ = collect(Float64, Array(u) + 0.2 * randn(size(u))) - time = vec(collect(Float64, ta)) - dataset = [x̂, time] - physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] - - # testing points for solve call(saveat=1/50.0 ∴ at t = collect(eltype(saveat), prob.tspan[1]:saveat:prob.tspan[2] internally estimates) - ta0 = range(tspan[1], tspan[2], length = 101) - u1 = [linear_analytic(u0, p, ti) for ti in ta0] - x̂1 = collect(Float64, Array(u1) + 0.2 * randn(size(u1))) - time1 = vec(collect(Float64, ta0)) - physsol1_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] - - chainlux1 = Chain(Dense(1, 7, tanh), Dense(7, 1)) - θinit, st = Lux.setup(Random.default_rng(), chainlux1) - - fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( - prob, chainlux1, dataset = dataset, draw_samples = 2500, - physdt = 1 / 50.0, priorsNNw = (0.0, 3.0), param = [LogNormal(9, 0.5)]) - - alg = BNNODE(chainlux1, dataset = dataset, draw_samples = 2500, physdt = 1 / 50.0, - priorsNNw = (0.0, 3.0), param = [LogNormal(9, 0.5)]) - - sol2lux = solve(prob, alg) - - # testing points - t = time - # Mean of last 500 sampled parameter's curves(flux and lux chains)[Ensemble predictions] - θ = [vector_to_parameters(fhsamples[i][1:(end - 1)], θinit) - for i in 2000:length(fhsamples)] - luxar = [chainlux1(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - # --------------------- ahmc_bayesian_pinn_ode() call - @test mean(abs.(physsol1 .- meanscurve)) < 0.15 - - # ESTIMATED ODE PARAMETERS (NN1 AND NN2) - @test abs(p - mean([fhsamples[i][23] for i in 2000:length(fhsamples)])) < abs(0.35 * p) - - #-------------------------- solve() call - @test mean(abs.(physsol1_1 .- pmean(sol2lux.ensemblesol[1]))) < 8e-2 - - # ESTIMATED ODE PARAMETERS (NN1 AND NN2) - @test abs(p - sol2lux.estimated_de_params[1]) < abs(0.15 * p) +@testitem "BPINN ODE II: With Parameter Estimation" tags = [:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + linear = (u, p, t) -> cos(p * t) + tspan = (0.0, 2.0) + u0 = 0.0 + p = 2 * pi + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, p) + + # Numerical and Analytical Solutions + sol1 = solve(prob, Tsit5(); saveat = 0.01) + u = sol1.u + time = sol1.t + + # BPINN AND TRAINING DATASET CREATION(dataset must be defined only inside problem timespan!) + ta = range(tspan[1], tspan[2], length = 100) + u = [linear_analytic(u0, p, ti) for ti in ta] + x̂ = collect(Float64, Array(u) + 0.2 * randn(size(u))) + time = vec(collect(Float64, ta)) + dataset = [x̂, time] + physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] + + # testing points for solve call(saveat=1/50.0 ∴ at t = collect(eltype(saveat), prob.tspan[1]:saveat:prob.tspan[2] internally estimates) + ta0 = range(tspan[1], tspan[2], length = 101) + u1 = [linear_analytic(u0, p, ti) for ti in ta0] + x̂1 = collect(Float64, Array(u1) + 0.2 * randn(size(u1))) + time1 = vec(collect(Float64, ta0)) + physsol1_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] + + chainlux1 = Chain(Dense(1, 7, tanh), Dense(7, 1)) + θinit, st = Lux.setup(Random.default_rng(), chainlux1) + + fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( + prob, chainlux1, dataset = dataset, draw_samples = 2500, + physdt = 1 / 50.0, priorsNNw = (0.0, 3.0), param = [LogNormal(9, 0.5)]) + + alg = BNNODE(chainlux1, dataset = dataset, draw_samples = 2500, physdt = 1 / 50.0, + priorsNNw = (0.0, 3.0), param = [LogNormal(9, 0.5)]) + + sol2lux = solve(prob, alg) + + # testing points + t = time + # Mean of last 500 sampled parameter's curves(flux and lux chains)[Ensemble predictions] + θ = [vector_to_parameters(fhsamples[i][1:(end-1)], θinit) + for i in 2000:length(fhsamples)] + luxar = [chainlux1(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + # --------------------- ahmc_bayesian_pinn_ode() call + @test mean(abs.(physsol1 .- meanscurve)) < 0.15 + + # ESTIMATED ODE PARAMETERS (NN1 AND NN2) + @test abs(p - mean([fhsamples[i][23] for i in 2000:length(fhsamples)])) < abs(0.35 * p) + + #-------------------------- solve() call + @test mean(abs.(physsol1_1 .- pmean(sol2lux.ensemblesol[1]))) < 8e-2 + + # ESTIMATED ODE PARAMETERS (NN1 AND NN2) + @test abs(p - sol2lux.estimated_de_params[1]) < abs(0.15 * p) end -@testitem "BPINN ODE III" tags=[:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - import Flux - - Random.seed!(100) - - linear = (u, p, t) -> u / p + exp(t / p) * cos(t) - tspan = (0.0, 10.0) - u0 = 0.0 - p = -5.0 - prob = ODEProblem(linear, u0, tspan, p) - linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) - # SOLUTION AND CREATE DATASET - sol = solve(prob, Tsit5(); saveat = 0.1) - u = sol.u - time = sol.t - x̂ = u .+ (u .* 0.1) .* randn(size(u)) - dataset = [x̂, time] - physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] - - # separate set of points for testing the solve() call (it uses saveat 1/50 hence here length 501) - time1 = vec(collect(Float64, range(tspan[1], tspan[2], length = 501))) - physsol2 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] - - chainlux12 = Chain(Dense(1, 6, tanh), Dense(6, 6, tanh), Dense(6, 1)) - θinit, st = Lux.setup(Random.default_rng(), chainlux12) - - # this a forward solve - fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode( - prob, chainlux12, draw_samples = 500, phystd = [0.01], priorsNNw = (0.0, 10.0)) - - fh_mcmc_chainlux22, fhsampleslux22, fhstatslux22 = ahmc_bayesian_pinn_ode( - prob, chainlux12, dataset = dataset, draw_samples = 500, l2std = [0.02], - phystd = [0.05], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)]) - - alg = BNNODE(chainlux12, dataset = dataset, draw_samples = 500, l2std = [0.02], - phystd = [0.05], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)]) - - sol3lux_pestim = solve(prob, alg) - - # testing timepoints - t = sol.t - #------------------------------ ahmc_bayesian_pinn_ode() call - # Mean of last 500 sampled parameter's curves(lux chains)[Ensemble predictions] - θ = [vector_to_parameters(fhsampleslux12[i], θinit) - for i in 400:length(fhsampleslux12)] - luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve2_1 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - θ = [vector_to_parameters(fhsampleslux22[i][1:(end - 1)], θinit) - for i in 400:length(fhsampleslux22)] - luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve2_2 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - @test mean(abs, sol.u .- meanscurve2_1) < 1e-2 - @test mean(abs, physsol1 .- meanscurve2_1) < 1e-2 - @test mean(abs, sol.u .- meanscurve2_2) < 1.5 - @test mean(abs, physsol1 .- meanscurve2_2) < 1.5 - - # estimated parameters(lux chain) - param1 = mean(i[62] for i in fhsampleslux22[400:length(fhsampleslux22)]) - @test abs(param1 - p) < abs(0.5 * p) +@testitem "BPINN ODE III" tags = [:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear = (u, p, t) -> u / p + exp(t / p) * cos(t) + tspan = (0.0, 10.0) + u0 = 0.0 + p = -5.0 + prob = ODEProblem(linear, u0, tspan, p) + linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) + # SOLUTION AND CREATE DATASET + sol = solve(prob, Tsit5(); saveat = 0.1) + u = sol.u + time = sol.t + x̂ = u .+ (u .* 0.1) .* randn(size(u)) + dataset = [x̂, time] + physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] + + # separate set of points for testing the solve() call (it uses saveat 1/50 hence here length 501) + time1 = vec(collect(Float64, range(tspan[1], tspan[2], length = 501))) + physsol2 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] + + chainlux12 = Chain(Dense(1, 6, tanh), Dense(6, 6, tanh), Dense(6, 1)) + θinit, st = Lux.setup(Random.default_rng(), chainlux12) + + # this a forward solve + fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode( + prob, chainlux12, draw_samples = 500, phystd = [0.01], priorsNNw = (0.0, 10.0)) + + fh_mcmc_chainlux22, fhsampleslux22, fhstatslux22 = ahmc_bayesian_pinn_ode( + prob, chainlux12, dataset = dataset, draw_samples = 500, l2std = [0.02], + phystd = [0.05], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)]) + + alg = BNNODE(chainlux12, dataset = dataset, draw_samples = 500, l2std = [0.02], + phystd = [0.05], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)]) + + sol3lux_pestim = solve(prob, alg) + + # testing timepoints + t = sol.t + #------------------------------ ahmc_bayesian_pinn_ode() call + # Mean of last 500 sampled parameter's curves(lux chains)[Ensemble predictions] + θ = [vector_to_parameters(fhsampleslux12[i], θinit) + for i in 400:length(fhsampleslux12)] + luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve2_1 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + θ = [vector_to_parameters(fhsampleslux22[i][1:(end-1)], θinit) + for i in 400:length(fhsampleslux22)] + luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve2_2 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + @test mean(abs, sol.u .- meanscurve2_1) < 1e-2 + @test mean(abs, physsol1 .- meanscurve2_1) < 1e-2 + @test mean(abs, sol.u .- meanscurve2_2) < 1.5 + @test mean(abs, physsol1 .- meanscurve2_2) < 1.5 + + # estimated parameters(lux chain) + param1 = mean(i[62] for i in fhsampleslux22[400:length(fhsampleslux22)]) + @test abs(param1 - p) < abs(0.5 * p) end -@testitem "BPINN ODE: Translating from Flux" tags=[:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - import Flux - - Random.seed!(100) - - linear_analytic = (u0, p, t) -> u0 + sin(2 * π * t) / (2 * π) - linear = (u, p, t) -> cos(2 * π * t) - tspan = (0.0, 2.0) - u0 = 0.0 - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan) - p = prob.p - - # Numerical and Analytical Solutions: testing ahmc_bayesian_pinn_ode() - ta = range(tspan[1], tspan[2], length = 300) - u = [linear_analytic(u0, nothing, ti) for ti in ta] - x̂ = collect(Float64, Array(u) + 0.02 * randn(size(u))) - time = vec(collect(Float64, ta)) - physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] - - # testing points for solve() call must match saveat(1/50.0) arg - ta0 = range(tspan[1], tspan[2], length = 101) - u1 = [linear_analytic(u0, nothing, ti) for ti in ta0] - x̂1 = collect(Float64, Array(u1) + 0.02 * randn(size(u1))) - time1 = vec(collect(Float64, ta0)) - physsol0_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] - chainflux = Flux.Chain(Flux.Dense(1, 7, tanh), Flux.Dense(7, 1)) |> Flux.f64 - fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( - prob, chainflux, draw_samples = 2500) - alg = BNNODE(chainflux, draw_samples = 2500) - @test alg.chain isa AbstractLuxLayer +@testitem "BPINN ODE: Translating from Flux" tags = [:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear_analytic = (u0, p, t) -> u0 + sin(2 * π * t) / (2 * π) + linear = (u, p, t) -> cos(2 * π * t) + tspan = (0.0, 2.0) + u0 = 0.0 + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan) + p = prob.p + + # Numerical and Analytical Solutions: testing ahmc_bayesian_pinn_ode() + ta = range(tspan[1], tspan[2], length = 300) + u = [linear_analytic(u0, nothing, ti) for ti in ta] + x̂ = collect(Float64, Array(u) + 0.02 * randn(size(u))) + time = vec(collect(Float64, ta)) + physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] + + # testing points for solve() call must match saveat(1/50.0) arg + ta0 = range(tspan[1], tspan[2], length = 101) + u1 = [linear_analytic(u0, nothing, ti) for ti in ta0] + x̂1 = collect(Float64, Array(u1) + 0.02 * randn(size(u1))) + time1 = vec(collect(Float64, ta0)) + physsol0_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] + chainflux = Flux.Chain(Flux.Dense(1, 7, tanh), Flux.Dense(7, 1)) |> Flux.f64 + fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( + prob, chainflux, draw_samples = 2500) + alg = BNNODE(chainflux, draw_samples = 2500) + @test alg.chain isa AbstractLuxLayer end -@testitem "BPINN ODE III: with the new objective" tags=[:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - import Flux - - Random.seed!(100) - - linear = (u, p, t) -> u / p + exp(t / p) * cos(t) - tspan = (0.0, 10.0) - u0 = 0.0 - p = -5.0 - prob = ODEProblem(linear, u0, tspan, p) - linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) - - # SOLUTION AND CREATE DATASET - sol = solve(prob, Tsit5(); saveat = 0.1) - u = sol.u - time = sol.t - x̂ = u .+ (0.1 .* randn(size(u))) - dataset = [x̂, time] - physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] - - chainlux12 = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) - θinit, st = Lux.setup(Random.default_rng(), chainlux12) - - fh_mcmc_chainlux22, fhsampleslux22, fhstatslux22 = ahmc_bayesian_pinn_ode( - prob, chainlux12, - dataset = dataset, - draw_samples = 500, - l2std = [0.1], - phystd = [0.01], - phynewstd = [0.01], - priorsNNw = (0.0, - 1.0), - param = [ - Normal(-7, 3) - ], estim_collocate = true) - - fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode( - prob, chainlux12, - dataset = dataset, - draw_samples = 500, - l2std = [0.1], - phystd = [0.01], - priorsNNw = (0.0, - 1.0), - param = [ - Normal(-7, 3) - ]) - - # testing timepoints - t = sol.t - #------------------------------ ahmc_bayesian_pinn_ode() call - # Mean of last 100 sampled parameter's curves(lux chains)[Ensemble predictions] - θ = [vector_to_parameters(fhsampleslux12[i][1:(end - 1)], θinit) - for i in 400:length(fhsampleslux12)] - luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve2_1 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - θ = [vector_to_parameters(fhsampleslux22[i][1:(end - 1)], θinit) - for i in 400:length(fhsampleslux22)] - luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve2_2 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - @test mean(abs.(sol.u .- meanscurve2_2)) < 1e-2 - @test mean(abs.(physsol1 .- meanscurve2_2)) < 1e-2 - @test mean(abs.(sol.u .- meanscurve2_1)) > mean(abs.(sol.u .- meanscurve2_2)) - @test mean(abs.(physsol1 .- meanscurve2_1)) > mean(abs.(physsol1 .- meanscurve2_2)) - - # estimated parameters(lux chain) - param2 = mean(i[62] for i in fhsampleslux22[400:length(fhsampleslux22)]) - @test abs(param2 - p) < abs(0.05 * p) - - param1 = mean(i[62] for i in fhsampleslux12[400:length(fhsampleslux12)]) - @test abs(param1 - p) > abs(0.5 * p) - @test abs(param2 - p) < abs(param1 - p) +@testitem "BPINN ODE III: with the new objective" tags = [:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear = (u, p, t) -> u / p + exp(t / p) * cos(t) + tspan = (0.0, 10.0) + u0 = 0.0 + p = -5.0 + prob = ODEProblem(linear, u0, tspan, p) + linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) + + # SOLUTION AND CREATE DATASET + sol = solve(prob, Tsit5(); saveat = 0.1) + u = sol.u + time = sol.t + x̂ = u .+ (0.1 .* randn(size(u))) + dataset = [x̂, time] + physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] + + chainlux12 = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) + θinit, st = Lux.setup(Random.default_rng(), chainlux12) + + fh_mcmc_chainlux22, fhsampleslux22, fhstatslux22 = ahmc_bayesian_pinn_ode( + prob, chainlux12, + dataset = dataset, + draw_samples = 500, + l2std = [0.1], + phystd = [0.01], + phynewstd = [0.01], + priorsNNw = (0.0, + 1.0), + param = [ + Normal(-7, 3), + ], estim_collocate = true) + + fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode( + prob, chainlux12, + dataset = dataset, + draw_samples = 500, + l2std = [0.1], + phystd = [0.01], + priorsNNw = (0.0, + 1.0), + param = [ + Normal(-7, 3), + ]) + + # testing timepoints + t = sol.t + #------------------------------ ahmc_bayesian_pinn_ode() call + # Mean of last 100 sampled parameter's curves(lux chains)[Ensemble predictions] + θ = [vector_to_parameters(fhsampleslux12[i][1:(end-1)], θinit) + for i in 400:length(fhsampleslux12)] + luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve2_1 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + θ = [vector_to_parameters(fhsampleslux22[i][1:(end-1)], θinit) + for i in 400:length(fhsampleslux22)] + luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve2_2 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + @test mean(abs.(sol.u .- meanscurve2_2)) < 1e-2 + @test mean(abs.(physsol1 .- meanscurve2_2)) < 1e-2 + @test mean(abs.(sol.u .- meanscurve2_1)) > mean(abs.(sol.u .- meanscurve2_2)) + @test mean(abs.(physsol1 .- meanscurve2_1)) > mean(abs.(physsol1 .- meanscurve2_2)) + + # estimated parameters(lux chain) + param2 = mean(i[62] for i in fhsampleslux22[400:length(fhsampleslux22)]) + @test abs(param2 - p) < abs(0.05 * p) + + param1 = mean(i[62] for i in fhsampleslux12[400:length(fhsampleslux12)]) + @test abs(param1 - p) > abs(0.5 * p) + @test abs(param2 - p) < abs(param1 - p) end -@testitem "BPINN ODE III: new objective solve call" tags=[:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - import Flux - - Random.seed!(100) - - linear = (u, p, t) -> u / p + exp(t / p) * cos(t) - tspan = (0.0, 10.0) - u0 = 0.0 - p = -5.0 - prob = ODEProblem(linear, u0, tspan, p) - linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) - - # SOLUTION AND CREATE DATASET - sol = solve(prob, Tsit5(); saveat = 0.1) - u = sol.u - time = sol.t - x̂ = u .+ (0.1 .* randn(size(u))) - dataset = [x̂, time] - - # set of points for testing the solve() call (it uses saveat 1/50 hence here length 501) - time1 = vec(collect(Float64, range(tspan[1], tspan[2], length = 501))) - physsol2 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] - - chainlux12 = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) - θinit, st = Lux.setup(Random.default_rng(), chainlux12) - - alg = BNNODE(chainlux12, - dataset = dataset, - draw_samples = 1000, - l2std = [0.1], - phystd = [0.01], - phynewstd = [0.01], - priorsNNw = (0.0, - 1.0), - param = [ - Normal(-7, 3) - ], numensemble = 200, - estim_collocate = true) - - sol3lux_pestim = solve(prob, alg) - - #-------------------------- solve() call - @test mean(abs.(physsol2 .- pmean(sol3lux_pestim.ensemblesol[1]))) < 1e-2 - - # estimated parameters - param3 = sol3lux_pestim.estimated_de_params[1] - @test abs(param3 - p) < abs(0.05 * p) +@testitem "BPINN ODE III: new objective solve call" tags = [:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear = (u, p, t) -> u / p + exp(t / p) * cos(t) + tspan = (0.0, 10.0) + u0 = 0.0 + p = -5.0 + prob = ODEProblem(linear, u0, tspan, p) + linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) + + # SOLUTION AND CREATE DATASET + sol = solve(prob, Tsit5(); saveat = 0.1) + u = sol.u + time = sol.t + x̂ = u .+ (0.1 .* randn(size(u))) + dataset = [x̂, time] + + # set of points for testing the solve() call (it uses saveat 1/50 hence here length 501) + time1 = vec(collect(Float64, range(tspan[1], tspan[2], length = 501))) + physsol2 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] + + chainlux12 = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) + θinit, st = Lux.setup(Random.default_rng(), chainlux12) + + alg = BNNODE(chainlux12, + dataset = dataset, + draw_samples = 1000, + l2std = [0.1], + phystd = [0.01], + phynewstd = [0.01], + priorsNNw = (0.0, + 1.0), + param = [ + Normal(-7, 3), + ], numensemble = 200, + estim_collocate = true) + + sol3lux_pestim = solve(prob, alg) + + #-------------------------- solve() call + @test mean(abs.(physsol2 .- pmean(sol3lux_pestim.ensemblesol[1]))) < 1e-2 + + # estimated parameters + param3 = sol3lux_pestim.estimated_de_params[1] + @test abs(param3 - p) < abs(0.05 * p) end -@testitem "BPINN ODE IV: Improvement" tags=[:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - import Flux - - Random.seed!(100) - - function lotka_volterra(u, p, t) - # Model parameters. - α, δ = p - # Current state. - x, y = u - - # Evaluate differential equations. - dx = (α - y) * x # prey - dy = (x - δ) * y # predator - - return [dx, dy] - end - - # initial-value problem. - u0 = [1.0, 1.0] - p = [1.5, 3.0] - tspan = (0.0, 7.0) - prob = ODEProblem(lotka_volterra, u0, tspan, p) - - # OrdinaryDiffEq.jl solve - dt = 0.1 - solution = solve(prob, Tsit5(); saveat = dt) - - times = solution.t - u = hcat(solution.u...) - x = u[1, :] + (0.5 .* randn(length(u[1, :]))) - y = u[2, :] + (0.5 .* randn(length(u[2, :]))) - dataset = [x, y, times] - - chain = Lux.Chain(Lux.Dense(1, 7, tanh), Lux.Dense(7, 7, tanh), - Lux.Dense(7, 2)) - - alg1 = BNNODE(chain; - dataset = dataset, - draw_samples = 1000, - l2std = [0.5, 0.5], - phystd = [0.5, 0.5], - priorsNNw = (0.0, 1.0), - param = [ - Normal(2, 2), - Normal(2, 2)]) - - alg2 = BNNODE(chain; - dataset = dataset, - draw_samples = 1000, - l2std = [0.5, 0.5], - phystd = [0.5, 0.5], - phynewstd = [1.0, 1.0], - priorsNNw = (0.0, 1.0), - param = [ - Normal(2, 2), - Normal(2, 2)], estim_collocate = true) - - @time sol_pestim1 = solve(prob, alg1; saveat = dt) - @time sol_pestim2 = solve(prob, alg2; saveat = dt) - - unsafe_comparisons(true) - bitvec = abs.(p .- sol_pestim1.estimated_de_params) .> - abs.(p .- sol_pestim2.estimated_de_params) - @test bitvec == ones(size(bitvec)) - - Loss_1 = mean(abs, u[1, :] .- pmean(sol_pestim1.ensemblesol[1])) + - mean(abs, u[2, :] .- pmean(sol_pestim1.ensemblesol[2])) - Loss_2 = mean(abs, u[1, :] .- pmean(sol_pestim2.ensemblesol[1])) + - mean(abs, u[2, :] .- pmean(sol_pestim2.ensemblesol[2])) - - @test Loss_1 > Loss_2 -end \ No newline at end of file +@testitem "BPINN ODE IV: Improvement" tags = [:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + function lotka_volterra(u, p, t) + # Model parameters. + α, δ = p + # Current state. + x, y = u + + # Evaluate differential equations. + dx = (α - y) * x # prey + dy = (x - δ) * y # predator + + return [dx, dy] + end + + # initial-value problem. + u0 = [1.0, 1.0] + p = [1.5, 3.0] + tspan = (0.0, 7.0) + prob = ODEProblem(lotka_volterra, u0, tspan, p) + + # OrdinaryDiffEq.jl solve + dt = 0.1 + solution = solve(prob, Tsit5(); saveat = dt) + + times = solution.t + u = hcat(solution.u...) + x = u[1, :] + (0.5 .* randn(length(u[1, :]))) + y = u[2, :] + (0.5 .* randn(length(u[2, :]))) + dataset = [x, y, times] + + chain = Lux.Chain(Lux.Dense(1, 7, tanh), Lux.Dense(7, 7, tanh), + Lux.Dense(7, 2)) + + alg1 = BNNODE(chain; + dataset = dataset, + draw_samples = 1000, + l2std = [0.5, 0.5], + phystd = [0.5, 0.5], + priorsNNw = (0.0, 1.0), + param = [ + Normal(2, 2), + Normal(2, 2)]) + + alg2 = BNNODE(chain; + dataset = dataset, + draw_samples = 1000, + l2std = [0.5, 0.5], + phystd = [0.5, 0.5], + phynewstd = [1.0, 1.0], + priorsNNw = (0.0, 1.0), + param = [ + Normal(2, 2), + Normal(2, 2)], estim_collocate = true) + + @time sol_pestim1 = solve(prob, alg1; saveat = dt) + @time sol_pestim2 = solve(prob, alg2; saveat = dt) + + unsafe_comparisons(true) + bitvec = abs.(p .- sol_pestim1.estimated_de_params) .> + abs.(p .- sol_pestim2.estimated_de_params) + @test bitvec == ones(size(bitvec)) + + Loss_1 = mean(abs, u[1, :] .- pmean(sol_pestim1.ensemblesol[1])) + + mean(abs, u[2, :] .- pmean(sol_pestim1.ensemblesol[2])) + Loss_2 = mean(abs, u[1, :] .- pmean(sol_pestim2.ensemblesol[1])) + + mean(abs, u[2, :] .- pmean(sol_pestim2.ensemblesol[2])) + + @test Loss_1 > Loss_2 +end diff --git a/lib/BayesianNeuralPDE/test/runtests.jl b/lib/BayesianNeuralPDE/test/runtests.jl index d84f163db..75d7889d1 100644 --- a/lib/BayesianNeuralPDE/test/runtests.jl +++ b/lib/BayesianNeuralPDE/test/runtests.jl @@ -13,8 +13,8 @@ const GROUP = lowercase(get(ENV, "GROUP", "all")) using NeuralPDE @info "Running tests with $(RETESTITEMS_NWORKERS) workers and \ - $(RETESTITEMS_NWORKER_THREADS) threads for group $(GROUP)" + $(RETESTITEMS_NWORKER_THREADS) threads for group $(GROUP)" ReTestItems.runtests(NeuralPDE; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]), - nworkers = RETESTITEMS_NWORKERS, - nworker_threads = RETESTITEMS_NWORKER_THREADS, testitem_timeout = 3600) \ No newline at end of file + nworkers = RETESTITEMS_NWORKERS, + nworker_threads = RETESTITEMS_NWORKER_THREADS, testitem_timeout = 3600) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index ba947fa7e..fd5214eda 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -28,8 +28,8 @@ using RecursiveArrayTools: DiffEqArray using Reexport: @reexport using RuntimeGeneratedFunctions: RuntimeGeneratedFunctions, @RuntimeGeneratedFunction using SciMLBase: SciMLBase, BatchIntegralFunction, IntegralProblem, NoiseProblem, - OptimizationFunction, OptimizationProblem, ReturnCode, discretize, - isinplace, solve, symbolic_discretize + OptimizationFunction, OptimizationProblem, ReturnCode, discretize, + isinplace, solve, symbolic_discretize using Statistics: Statistics, mean using QuasiMonteCarlo: QuasiMonteCarlo, LatinHypercubeSample using WeightInitializers: glorot_uniform, zeros32 @@ -38,18 +38,16 @@ using Zygote: Zygote # Symbolic Stuff using ModelingToolkit: ModelingToolkit, PDESystem, Differential, toexpr using Symbolics: Symbolics, unwrap, arguments, operation, build_expr, Num, - expand_derivatives + expand_derivatives using SymbolicUtils: SymbolicUtils using SymbolicIndexingInterface: SymbolicIndexingInterface # Needed for the Bayesian Stuff -using AdvancedHMC: AdvancedHMC, DiagEuclideanMetric, HMC, HMCDA, Hamiltonian, - JitteredLeapfrog, Leapfrog, MassMatrixAdaptor, NUTS, StanHMCAdaptor, - StepSizeAdaptor, TemperedLeapfrog, find_good_stepsize -using Distributions: Distributions, Distribution, MvNormal, Normal, dim, logpdf +using AdvancedHMC: AdvancedHMC, HMCDA, + NUTS +using Distributions: Distributions, Distribution, MvNormal, dim, logpdf using LogDensityProblems: LogDensityProblems -using MCMCChains: MCMCChains, Chains, sample -using MonteCarloMeasurements: Particles +using MCMCChains: MCMCChains import LuxCore: initialparameters, initialstates, parameterlength @@ -97,19 +95,19 @@ export DeepGalerkin using .BayesianNeuralPDE export ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde, - BNNODE, BPINNsolution, BayesianPINN + BNNODE, BPINNsolution, BayesianPINN export neural_adapter export GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, - WeightedIntervalTraining + WeightedIntervalTraining export build_loss_function, get_loss_function, - generate_training_sets, get_variables, get_argument, get_bounds, - get_numeric_integral, symbolic_discretize, vector_to_parameters + generate_training_sets, get_variables, get_argument, get_bounds, + get_numeric_integral, symbolic_discretize, vector_to_parameters export AbstractAdaptiveLoss, NonAdaptiveLoss, GradientScaleAdaptiveLoss, - MiniMaxAdaptiveLoss + MiniMaxAdaptiveLoss export LogOptions From c7c26f15fbfadae2aa0dabfab42668cd08ef5c0f Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sun, 2 Feb 2025 09:56:35 +0530 Subject: [PATCH 13/15] Removed BayesianNeuralPDE imports --- src/NeuralPDE.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index fd5214eda..3a31bf1af 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -75,7 +75,6 @@ include("pinn_types.jl") include("symbolic_utilities.jl") include("training_strategies.jl") include("adaptive_losses.jl") -include("../lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl") include("ode_solve.jl") include("dae_solve.jl") @@ -93,7 +92,6 @@ export NNODE, NNDAE export PhysicsInformedNN, discretize export DeepGalerkin -using .BayesianNeuralPDE export ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde, BNNODE, BPINNsolution, BayesianPINN From dffd8226121feb9f236ff94a8b298cd2f8d005f5 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sun, 2 Feb 2025 15:15:55 +0530 Subject: [PATCH 14/15] Removed undefined exports --- src/NeuralPDE.jl | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index 3a31bf1af..bcd78c170 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -92,9 +92,6 @@ export NNODE, NNDAE export PhysicsInformedNN, discretize export DeepGalerkin -export ahmc_bayesian_pinn_ode, ahmc_bayesian_pinn_pde, - BNNODE, BPINNsolution, BayesianPINN - export neural_adapter export GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, From 57ee1aa865a25c5e4a668c6fb26f480f7cf65393 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sun, 2 Feb 2025 15:22:49 +0530 Subject: [PATCH 15/15] Format --- lib/BayesianNeuralPDE/src/BPINN_ode.jl | 226 ++--- .../src/BayesianNeuralPDE.jl | 2 +- lib/BayesianNeuralPDE/src/PDE_BPINN.jl | 821 +++++++-------- lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl | 268 ++--- lib/BayesianNeuralPDE/src/discretize.jl | 122 +-- lib/BayesianNeuralPDE/src/pinn_types.jl | 14 +- lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl | 936 +++++++++--------- lib/BayesianNeuralPDE/test/BPINN_tests.jl | 832 ++++++++-------- lib/BayesianNeuralPDE/test/runtests.jl | 6 +- src/NeuralPDE.jl | 16 +- 10 files changed, 1622 insertions(+), 1621 deletions(-) diff --git a/lib/BayesianNeuralPDE/src/BPINN_ode.jl b/lib/BayesianNeuralPDE/src/BPINN_ode.jl index 68d996ba5..0b3589742 100644 --- a/lib/BayesianNeuralPDE/src/BPINN_ode.jl +++ b/lib/BayesianNeuralPDE/src/BPINN_ode.jl @@ -78,42 +78,42 @@ Kevin Linka, Amelie Schäfer, Xuhui Meng, Zongren Zou, George Em Karniadakis, El "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 + 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) + 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 """ @@ -132,9 +132,9 @@ Contains `ahmc_bayesian_pinn_ode()` function output: - nom_step_size """ @concrete struct BPINNstats - mcmc_chain::Any - samples::Any - statistics::Any + mcmc_chain::Any + samples::Any + statistics::Any end """ @@ -148,84 +148,84 @@ contains fields related to that). parameters. """ @concrete struct BPINNsolution - original <: BPINNstats - ensemblesol::Any - estimated_nn_params::Any - estimated_de_params::Any - timepoints::Any + 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) + 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 diff --git a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl index 38aaa696f..3944a8870 100644 --- a/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl +++ b/lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl @@ -1,7 +1,7 @@ module BayesianNeuralPDE using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements using Printf: @printf using ConcreteStructs: @concrete using NeuralPDE: PhysicsInformedNN diff --git a/lib/BayesianNeuralPDE/src/PDE_BPINN.jl b/lib/BayesianNeuralPDE/src/PDE_BPINN.jl index beffe557e..2b55ac305 100644 --- a/lib/BayesianNeuralPDE/src/PDE_BPINN.jl +++ b/lib/BayesianNeuralPDE/src/PDE_BPINN.jl @@ -1,253 +1,254 @@ @concrete struct PDELogTargetDensity - dim::Int - strategy <: AbstractTrainingStrategy - dataset <: Union{Nothing, Vector{<:Matrix{<:Real}}} - priors <: Vector{<:Distribution} - allstd::Vector{Vector{Float64}} - phynewstd::Vector{Float64} - names::Tuple - extraparams::Int - init_params <: Union{AbstractVector, NamedTuple, ComponentArray} - full_loglikelihood::Any - L2_loss2::Any - Φ::Any + dim::Int + strategy <: AbstractTrainingStrategy + dataset <: Union{Nothing, Vector{<:Matrix{<:Real}}} + priors <: Vector{<:Distribution} + allstd::Vector{Vector{Float64}} + phynewstd::Vector{Float64} + names::Tuple + extraparams::Int + init_params <: Union{AbstractVector, NamedTuple, ComponentArray} + full_loglikelihood::Any + L2_loss2::Any + Φ::Any end function LogDensityProblems.logdensity(ltd::PDELogTargetDensity, θ) - # for parameter estimation neccesarry to use multioutput case - if ltd.L2_loss2 === nothing - return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) + - priorlogpdf(ltd, θ) + L2LossData(ltd, θ) - else - return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) + - priorlogpdf(ltd, θ) + L2LossData(ltd, θ) + ltd.L2_loss2(setparameters(ltd, θ), ltd.phynewstd) - end + # for parameter estimation neccesarry to use multioutput case + if ltd.L2_loss2 === nothing + return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) + + priorlogpdf(ltd, θ) + L2LossData(ltd, θ) + else + return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) + + priorlogpdf(ltd, θ) + L2LossData(ltd, θ) + + ltd.L2_loss2(setparameters(ltd, θ), ltd.phynewstd) + end end # you get a vector of losses function get_lossy(pinnrep, dataset, Dict_differentials) - eqs = pinnrep.eqs - depvars = pinnrep.depvars #depvar order is same as dataset - - # Dict_differentials is filled with Differential operator => diff_i key-value pairs - # masking operation - eqs_new = SymbolicUtils.substitute.(eqs, Ref(Dict_differentials)) - - to_subs, tobe_subs = get_symbols(dataset, depvars, eqs) - - # for values of all depvars at corresponding indvar values in dataset, create dictionaries {Dict(x(t) => 1.0496435863173237, y(t) => 1.9227770685615337)} - # In each Dict, num form of depvar is key to its value at certain coords of indvars, n_dicts = n_rows_dataset(or n_indvar_coords_dataset) - eq_subs = [Dict(tobe_subs[depvar] => to_subs[depvar][i] for depvar in depvars) - for i in 1:size(dataset[1][:, 1])[1]] - - # for each dataset point(eq_sub dictionary), substitute in masked equations - # n_collocated_equations = n_rows_dataset(or n_indvar_coords_dataset) - masked_colloc_equations = [[Symbolics.substitute(eq, eq_sub) for eq in eqs_new] - for eq_sub in eq_subs] - # now we have vector of dataset depvar's collocated equations - - # reverse dict for re-substituting values of Differential(t)(u(t)) etc - rev_Dict_differentials = Dict(value => key for (key, value) in Dict_differentials) - - # unmask Differential terms in masked_colloc_equations - colloc_equations = [Symbolics.substitute.( - masked_colloc_equation, Ref(rev_Dict_differentials)) - for masked_colloc_equation in masked_colloc_equations] - # nested vector of data_pde_loss_functions (as in discretize.jl) - # each sub vector has dataset's indvar coord's datafree_colloc_loss_function, n_subvectors = n_rows_dataset(or n_indvar_coords_dataset) - # zip each colloc equation with args for each build_loss call per equation vector - data_colloc_loss_functions = [[build_loss_function(pinnrep, eq, pde_indvar) - for (eq, pde_indvar, integration_indvar) in zip( - colloc_equation, - pinnrep.pde_indvars, - pinnrep.pde_integration_vars)] - for colloc_equation in colloc_equations] - - return data_colloc_loss_functions + eqs = pinnrep.eqs + depvars = pinnrep.depvars #depvar order is same as dataset + + # Dict_differentials is filled with Differential operator => diff_i key-value pairs + # masking operation + eqs_new = SymbolicUtils.substitute.(eqs, Ref(Dict_differentials)) + + to_subs, tobe_subs = get_symbols(dataset, depvars, eqs) + + # for values of all depvars at corresponding indvar values in dataset, create dictionaries {Dict(x(t) => 1.0496435863173237, y(t) => 1.9227770685615337)} + # In each Dict, num form of depvar is key to its value at certain coords of indvars, n_dicts = n_rows_dataset(or n_indvar_coords_dataset) + eq_subs = [Dict(tobe_subs[depvar] => to_subs[depvar][i] for depvar in depvars) + for i in 1:size(dataset[1][:, 1])[1]] + + # for each dataset point(eq_sub dictionary), substitute in masked equations + # n_collocated_equations = n_rows_dataset(or n_indvar_coords_dataset) + masked_colloc_equations = [[Symbolics.substitute(eq, eq_sub) for eq in eqs_new] + for eq_sub in eq_subs] + # now we have vector of dataset depvar's collocated equations + + # reverse dict for re-substituting values of Differential(t)(u(t)) etc + rev_Dict_differentials = Dict(value => key for (key, value) in Dict_differentials) + + # unmask Differential terms in masked_colloc_equations + colloc_equations = [Symbolics.substitute.( + masked_colloc_equation, Ref(rev_Dict_differentials)) + for masked_colloc_equation in masked_colloc_equations] + # nested vector of data_pde_loss_functions (as in discretize.jl) + # each sub vector has dataset's indvar coord's datafree_colloc_loss_function, n_subvectors = n_rows_dataset(or n_indvar_coords_dataset) + # zip each colloc equation with args for each build_loss call per equation vector + data_colloc_loss_functions = [[build_loss_function(pinnrep, eq, pde_indvar) + for (eq, pde_indvar, integration_indvar) in zip( + colloc_equation, + pinnrep.pde_indvars, + pinnrep.pde_integration_vars)] + for colloc_equation in colloc_equations] + + return data_colloc_loss_functions end function get_symbols(dataset, depvars, eqs) - # take only values of depvars from dataset - depvar_vals = [dataset_i[:, 1] for dataset_i in dataset] - # order of pinnrep.depvars, depvar_vals, BayesianPINN.dataset must be same - to_subs = Dict(depvars .=> depvar_vals) - - numform_vars = Symbolics.get_variables.(eqs) - Eq_vars = unique(reduce(vcat, numform_vars)) - # got equation's depvar num format {x(t)} for use in substitute() - - tobe_subs = Dict() - for a in depvars - for i in Eq_vars - expr = toexpr(i) - if (expr isa Expr) && (expr.args[1] == a) - tobe_subs[a] = i - end - end - end - # depvar symbolic and num format got, tobe_subs : Dict{Any, Any}(:y => y(t), :x => x(t)) - - return to_subs, tobe_subs + # take only values of depvars from dataset + depvar_vals = [dataset_i[:, 1] for dataset_i in dataset] + # order of pinnrep.depvars, depvar_vals, BayesianPINN.dataset must be same + to_subs = Dict(depvars .=> depvar_vals) + + numform_vars = Symbolics.get_variables.(eqs) + Eq_vars = unique(reduce(vcat, numform_vars)) + # got equation's depvar num format {x(t)} for use in substitute() + + tobe_subs = Dict() + for a in depvars + for i in Eq_vars + expr = toexpr(i) + if (expr isa Expr) && (expr.args[1] == a) + tobe_subs[a] = i + end + end + end + # depvar symbolic and num format got, tobe_subs : Dict{Any, Any}(:y => y(t), :x => x(t)) + + return to_subs, tobe_subs end @views function setparameters(ltd::PDELogTargetDensity, θ) - names = ltd.names - ps_new = θ[1:(end-ltd.extraparams)] - ps = ltd.init_params - - # multioutput case for Lux chains, for each depvar ps would contain Lux ComponentVectors - # which we use for mapping current ahmc sampled vector of parameters onto NNs - i = 0 - Luxparams = [vector_to_parameters(ps_new[((i+=length(ps[x]))-length(ps[x])+1):i], - ps[x]) for x in names] - - a = ComponentArray(NamedTuple{ltd.names}(i for i in Luxparams)) - - if ltd.extraparams > 0 - return ComponentArray(; depvar = a, p = θ[(end-ltd.extraparams+1):end]) - else - return ComponentArray(; depvar = a) - end + names = ltd.names + ps_new = θ[1:(end - ltd.extraparams)] + ps = ltd.init_params + + # multioutput case for Lux chains, for each depvar ps would contain Lux ComponentVectors + # which we use for mapping current ahmc sampled vector of parameters onto NNs + i = 0 + Luxparams = [vector_to_parameters(ps_new[((i += length(ps[x])) - length(ps[x]) + 1):i], + ps[x]) for x in names] + + a = ComponentArray(NamedTuple{ltd.names}(i for i in Luxparams)) + + if ltd.extraparams > 0 + return ComponentArray(; depvar = a, p = θ[(end - ltd.extraparams + 1):end]) + else + return ComponentArray(; depvar = a) + end end LogDensityProblems.dimension(ltd::PDELogTargetDensity) = ltd.dim function LogDensityProblems.capabilities(::PDELogTargetDensity) - LogDensityProblems.LogDensityOrder{1}() + LogDensityProblems.LogDensityOrder{1}() end # L2 losses loglikelihood(needed mainly for ODE parameter estimation) function L2LossData(ltd::PDELogTargetDensity, θ) - Φ = ltd.Φ - init_params = ltd.init_params - dataset = ltd.dataset - L2stds = ltd.allstd[3] - # each dep var has a diff dataset depending on its indep var and their domains - # these datasets are matrices of first col-dep var and remaining cols-all indep var - # ltd.init_params is needed to construct a vector of parameters into a ComponentVector - - # dataset of form Vector[matrix_x, matrix_y, matrix_z] - # matrix_i is of form [i,indvar1,indvar2,..] (needed in case if heterogenous domains) - - # Phi is the trial solution for each NN in chain array - # Creating logpdf( MvNormal(Phi(t,θ),std), dataset[i] ) - # dataset[i][:, 2:end] -> indepvar cols of a particular depvar's dataset - # dataset[i][:, 1] -> depvar col of depvar's dataset - - ltd.extraparams ≤ 0 && return false - - sumt = 0 - for i in eachindex(Φ) - sumt += logpdf( - MvNormal( - Φ[i](dataset[i][:, 2:end]', - vector_to_parameters(θ[1:(end-ltd.extraparams)], init_params)[ltd.names[i]])[ - 1, :], - Diagonal(abs2.(ones(size(dataset[i])[1]) .* L2stds[i]))), - dataset[i][:, 1]) - end - return sumt + Φ = ltd.Φ + init_params = ltd.init_params + dataset = ltd.dataset + L2stds = ltd.allstd[3] + # each dep var has a diff dataset depending on its indep var and their domains + # these datasets are matrices of first col-dep var and remaining cols-all indep var + # ltd.init_params is needed to construct a vector of parameters into a ComponentVector + + # dataset of form Vector[matrix_x, matrix_y, matrix_z] + # matrix_i is of form [i,indvar1,indvar2,..] (needed in case if heterogenous domains) + + # Phi is the trial solution for each NN in chain array + # Creating logpdf( MvNormal(Phi(t,θ),std), dataset[i] ) + # dataset[i][:, 2:end] -> indepvar cols of a particular depvar's dataset + # dataset[i][:, 1] -> depvar col of depvar's dataset + + ltd.extraparams ≤ 0 && return false + + sumt = 0 + for i in eachindex(Φ) + sumt += logpdf( + MvNormal( + Φ[i](dataset[i][:, 2:end]', + vector_to_parameters(θ[1:(end - ltd.extraparams)], init_params)[ltd.names[i]])[ + 1, :], + Diagonal(abs2.(ones(size(dataset[i])[1]) .* L2stds[i]))), + dataset[i][:, 1]) + end + return sumt end # priors for NN parameters + ODE constants function priorlogpdf(ltd::PDELogTargetDensity, θ) - allparams = ltd.priors - # Vector of ode parameters priors - invpriors = allparams[2:end] - nnwparams = allparams[1] + allparams = ltd.priors + # Vector of ode parameters priors + invpriors = allparams[2:end] + nnwparams = allparams[1] - ltd.extraparams ≤ 0 && return logpdf(nnwparams, θ) + ltd.extraparams ≤ 0 && return logpdf(nnwparams, θ) - invlogpdf = sum((length(θ)-ltd.extraparams+1):length(θ)) do i - logpdf(invpriors[length(θ)-i+1], θ[i]) - end + invlogpdf = sum((length(θ) - ltd.extraparams + 1):length(θ)) do i + logpdf(invpriors[length(θ) - i + 1], θ[i]) + end - return invlogpdf + logpdf(nnwparams, θ[1:(length(θ)-ltd.extraparams)]) + return invlogpdf + logpdf(nnwparams, θ[1:(length(θ) - ltd.extraparams)]) end function integratorchoice(Integratorkwargs, initial_ϵ) - Integrator = Integratorkwargs[:Integrator] - if Integrator == JitteredLeapfrog - jitter_rate = Integratorkwargs[:jitter_rate] - Integrator(initial_ϵ, jitter_rate) - elseif Integrator == TemperedLeapfrog - tempering_rate = Integratorkwargs[:tempering_rate] - Integrator(initial_ϵ, tempering_rate) - else - Integrator(initial_ϵ) - end + Integrator = Integratorkwargs[:Integrator] + if Integrator == JitteredLeapfrog + jitter_rate = Integratorkwargs[:jitter_rate] + Integrator(initial_ϵ, jitter_rate) + elseif Integrator == TemperedLeapfrog + tempering_rate = Integratorkwargs[:tempering_rate] + Integrator(initial_ϵ, tempering_rate) + else + Integrator(initial_ϵ) + end end function adaptorchoice(Adaptor, mma, ssa) - if Adaptor != AdvancedHMC.NoAdaptation() - Adaptor(mma, ssa) - else - AdvancedHMC.NoAdaptation() - end + if Adaptor != AdvancedHMC.NoAdaptation() + Adaptor(mma, ssa) + else + AdvancedHMC.NoAdaptation() + end end function inference(samples, pinnrep, saveats, numensemble, ℓπ) - domains = pinnrep.domains - phi = pinnrep.phi - dict_depvar_input = pinnrep.dict_depvar_input - depvars = pinnrep.depvars - - names = ℓπ.names - initial_nnθ = ℓπ.init_params - ninv = ℓπ.extraparams - - ranges = Dict([Symbol(d.variables) => infimum(d.domain):dx:supremum(d.domain) - for (d, dx) in zip(domains, saveats)]) - inputs = [dict_depvar_input[i] for i in depvars] - - span = [[ranges[indvar] for indvar in input] for input in inputs] - timepoints = [hcat(vec(map(points -> collect(points), - Iterators.product(span[i]...)))...) - for i in eachindex(phi)] - - # order of range's domains must match chain's inputs and dep_vars - samples = samples[(end-numensemble):end] - nnparams = length(samples[1][1:(end-ninv)]) - # get rows-ith param and col-ith sample value - estimnnparams = [Particles(reduce(hcat, samples)[i, :]) - for i in 1:nnparams] - - # PDE params - if ninv == 0 - estimated_params = [nothing] - else - estimated_params = [Particles(reduce(hcat, samples)[i, :]) - for i in (nnparams+1):(nnparams+ninv)] - end - - # getting parameter ranges in case of Lux chains - Luxparams = [] - i = 0 - for x in names - len = length(initial_nnθ[x]) - push!(Luxparams, (i+1):(i+len)) - i += len - end - - # convert to format directly usable by lux - estimatedLuxparams = [vector_to_parameters(estimnnparams[Luxparams[i]], - initial_nnθ[names[i]]) for i in eachindex(phi)] - - # infer predictions(preds) each row - NN, each col - ith sample - samplesn = reduce(hcat, samples) - preds = [] - for j in eachindex(phi) - push!(preds, - [phi[j](timepoints[j], - vector_to_parameters(samplesn[:, i][Luxparams[j]], - initial_nnθ[names[j]])) for i in 1:numensemble]) - end - - # note here no of samples referse to numensemble and points is the no of points in each dep_vars discretization - # each phi will give output in single domain of depvar(so we have each row as a vector of vector outputs) - # so we get after reduce a single matrix of n rows(samples), and j cols(points) - ensemblecurves = [Particles(reduce(vcat, preds[i])) for i in eachindex(phi)] - return ensemblecurves, estimatedLuxparams, estimated_params, timepoints + domains = pinnrep.domains + phi = pinnrep.phi + dict_depvar_input = pinnrep.dict_depvar_input + depvars = pinnrep.depvars + + names = ℓπ.names + initial_nnθ = ℓπ.init_params + ninv = ℓπ.extraparams + + ranges = Dict([Symbol(d.variables) => infimum(d.domain):dx:supremum(d.domain) + for (d, dx) in zip(domains, saveats)]) + inputs = [dict_depvar_input[i] for i in depvars] + + span = [[ranges[indvar] for indvar in input] for input in inputs] + timepoints = [hcat(vec(map(points -> collect(points), + Iterators.product(span[i]...)))...) + for i in eachindex(phi)] + + # order of range's domains must match chain's inputs and dep_vars + samples = samples[(end - numensemble):end] + nnparams = length(samples[1][1:(end - ninv)]) + # get rows-ith param and col-ith sample value + estimnnparams = [Particles(reduce(hcat, samples)[i, :]) + for i in 1:nnparams] + + # PDE params + if ninv == 0 + estimated_params = [nothing] + else + estimated_params = [Particles(reduce(hcat, samples)[i, :]) + for i in (nnparams + 1):(nnparams + ninv)] + end + + # getting parameter ranges in case of Lux chains + Luxparams = [] + i = 0 + for x in names + len = length(initial_nnθ[x]) + push!(Luxparams, (i + 1):(i + len)) + i += len + end + + # convert to format directly usable by lux + estimatedLuxparams = [vector_to_parameters(estimnnparams[Luxparams[i]], + initial_nnθ[names[i]]) for i in eachindex(phi)] + + # infer predictions(preds) each row - NN, each col - ith sample + samplesn = reduce(hcat, samples) + preds = [] + for j in eachindex(phi) + push!(preds, + [phi[j](timepoints[j], + vector_to_parameters(samplesn[:, i][Luxparams[j]], + initial_nnθ[names[j]])) for i in 1:numensemble]) + end + + # note here no of samples referse to numensemble and points is the no of points in each dep_vars discretization + # each phi will give output in single domain of depvar(so we have each row as a vector of vector outputs) + # so we get after reduce a single matrix of n rows(samples), and j cols(points) + ensemblecurves = [Particles(reduce(vcat, preds[i])) for i in eachindex(phi)] + return ensemblecurves, estimatedLuxparams, estimated_params, timepoints end """ @@ -308,203 +309,203 @@ end releases. """ function 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), Dict_differentials = nothing, progress = false, verbose = false) - pinnrep = symbolic_discretize(pde_system, discretization) - dataset_pde, dataset_bc = discretization.dataset - - newloss = if Dict_differentials isa Nothing - nothing - else - data_colloc_loss_functions = get_lossy(pinnrep, dataset_pde, Dict_differentials) - # size = number of indvar coords in dataset - # add case for if parameters present in bcs? - - train_sets_pde = get_dataset_train_points(pde_system.eqs, - dataset_pde, - pinnrep) - # j is number of indvar coords in dataset, i is number of PDE equations in system - # -1 is placeholder, removed in merge_strategy_with_loglikelihood_function function call (train_sets[:, 2:end]()) - colloc_train_sets = [[hcat([-1], train_sets_pde[i][:, j]...) - for i in eachindex(data_colloc_loss_functions[1])] - for j in eachindex(data_colloc_loss_functions)] - - # using dataset's indvar coords as train_sets_pde and indvar coord's datafree_colloc_loss_function, create loss functions - # placeholder strategy = GridTraining(0.1), datafree_bc_loss_function and train_sets_bc must be nothing - # order of indvar coords will be same as corresponding depvar coords values in dataset provided in get_lossy() call. - pde_loss_function_points = [merge_strategy_with_loglikelihood_function( - pinnrep, - GridTraining(0.1), - data_colloc_loss_functions[i], - nothing; - train_sets_pde = colloc_train_sets[i], - train_sets_bc = nothing)[1] - for i in eachindex(data_colloc_loss_functions)] - - function L2_loss2(θ, phynewstd) - # first sum is over points losses over many equations for the same points - # second sum is over all points - pde_loglikelihoods = sum([sum([pde_loss_function(θ, - phynewstd[i]) - for (i, pde_loss_function) in enumerate(pde_loss_functions)]) - for pde_loss_functions in pde_loss_function_points]) - end - end - - # add overall functionality for BC dataset points (case of parametric BC) ? - if ((dataset_bc isa Nothing) && (dataset_pde isa Nothing)) - dataset = nothing - elseif dataset_bc isa Nothing - dataset = dataset_pde - elseif dataset_pde isa Nothing - dataset = dataset_bc - else - dataset = [vcat(dataset_pde[i], dataset_bc[i]) for i in eachindex(dataset_pde)] - end - - if discretization.param_estim && isempty(param) - throw(UndefVarError(:param)) - elseif discretization.param_estim && dataset isa Nothing - throw(UndefVarError(:dataset)) - elseif discretization.param_estim && length(l2std) != length(pinnrep.depvars) - error("L2 stds length must match number of dependant variables") - end - - # for physics loglikelihood - full_weighted_loglikelihood = pinnrep.loss_functions.full_loss_function - chain = discretization.chain - - if length(pinnrep.domains) != length(saveats) - error("Number of independent variables must match saveat inference discretization steps") - end - - # NN solutions for loglikelihood which is used for L2lossdata - Φ = pinnrep.phi - - @assert nchains ≥ 1 "number of chains must be greater than or equal to 1" - - # remove inv params take only NN params, AHMC uses Float64 - initial_nnθ = pinnrep.flat_init_params[1:(end-length(param))] - initial_θ = collect(Float64, initial_nnθ) - - # contains only NN parameters - initial_nnθ = pinnrep.init_params - - names = ntuple(i -> pinnrep.depvars[i], length(chain)) - - #ode parameter estimation - nparameters = length(initial_θ) - ninv = length(param) - # add init_params for NN params - priors = [ - MvNormal(priorsNNw[1] * ones(nparameters), - Diagonal(abs2.(priorsNNw[2] .* ones(nparameters)))), - ] - - # append Ode params to all paramvector - initial_θ - if ninv > 0 - # shift ode params(initialise ode params by prior means) - # check if means or user specified is better - initial_θ = vcat(initial_θ, [Distributions.params(param[i])[1] for i in 1:ninv]) - priors = vcat(priors, param) - nparameters += ninv - end - - # vector in case of N-dimensional domains - strategy = discretization.strategy - - # dimensions would be total no of params,initial_nnθ for Lux namedTuples - ℓπ = PDELogTargetDensity( - nparameters, strategy, dataset, priors, [phystd, bcstd, l2std], phynewstd, - names, ninv, initial_nnθ, full_weighted_loglikelihood, newloss, Φ) - - Adaptor, Metric, targetacceptancerate = Adaptorkwargs[:Adaptor], - Adaptorkwargs[:Metric], Adaptorkwargs[:targetacceptancerate] - - # Define Hamiltonian system (nparameters ~ dimensionality of the sampling space) - metric = Metric(nparameters) - hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) - - if verbose - @printf("Current Physics Log-likelihood : %g\n", - ℓπ.full_loglikelihood(setparameters(ℓπ, initial_θ), ℓπ.allstd)) - @printf("Current Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, initial_θ)) - @printf("Current SSE against dataset Log-likelihood : %g\n", - L2LossData(ℓπ, initial_θ)) - if !(newloss isa Nothing) - @printf("Current new loss : %g\n", - ℓπ.L2_loss2(setparameters(ℓπ, initial_θ), - ℓπ.phynewstd)) - end - end - - # parallel sampling option - if nchains != 1 - - # Cache to store the chains - bpinnsols = Vector{Any}(undef, nchains) - - Threads.@threads for i in 1:nchains - # each chain has different initial NNparameter values(better posterior exploration) - initial_θ = vcat(randn(nparameters - ninv), - initial_θ[(nparameters-ninv+1):end]) - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = integratorchoice(Integratorkwargs, initial_ϵ) - adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), - StepSizeAdaptor(targetacceptancerate, integrator)) - Kernel = AdvancedHMC.make_kernel(Kernel, integrator) - samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, adaptor; - progress = progress, verbose = verbose) - - # return a chain(basic chain),samples and stats - matrix_samples = hcat(samples...) - mcmc_chain = MCMCChains.Chains(matrix_samples') - - fullsolution = BPINNstats(mcmc_chain, samples, stats) - ensemblecurves, estimnnparams, estimated_params, timepoints = inference( - samples, pinnrep, saveat, numensemble, ℓπ) - - bpinnsols[i] = BPINNsolution( - fullsolution, ensemblecurves, estimnnparams, estimated_params, timepoints) - end - return bpinnsols - else - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = integratorchoice(Integratorkwargs, initial_ϵ) - adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), - StepSizeAdaptor(targetacceptancerate, integrator)) - - Kernel = AdvancedHMC.make_kernel(Kernel, integrator) - samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, - adaptor; progress = progress, verbose = verbose) - - # return a chain(basic chain),samples and stats - matrix_samples = hcat(samples...) - mcmc_chain = MCMCChains.Chains(matrix_samples') - - if verbose - @printf("Sampling Complete.\n") - @printf("Final Physics Log-likelihood : %g\n", - ℓπ.full_loglikelihood(setparameters(ℓπ, samples[end]), ℓπ.allstd)) - @printf("Final Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, samples[end])) - @printf("Final SSE against dataset Log-likelihood : %g\n", - L2LossData(ℓπ, samples[end])) - if !(newloss isa Nothing) - @printf("Final new loss : %g\n", - ℓπ.L2_loss2(setparameters(ℓπ, samples[end]), - ℓπ.phynewstd)) - end - end - - fullsolution = BPINNstats(mcmc_chain, samples, stats) - ensemblecurves, estimnnparams, estimated_params, timepoints = inference(samples, - pinnrep, saveats, numensemble, ℓπ) - - return BPINNsolution( - fullsolution, ensemblecurves, estimnnparams, estimated_params, timepoints) - end + 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), Dict_differentials = nothing, progress = false, verbose = false) + pinnrep = symbolic_discretize(pde_system, discretization) + dataset_pde, dataset_bc = discretization.dataset + + newloss = if Dict_differentials isa Nothing + nothing + else + data_colloc_loss_functions = get_lossy(pinnrep, dataset_pde, Dict_differentials) + # size = number of indvar coords in dataset + # add case for if parameters present in bcs? + + train_sets_pde = get_dataset_train_points(pde_system.eqs, + dataset_pde, + pinnrep) + # j is number of indvar coords in dataset, i is number of PDE equations in system + # -1 is placeholder, removed in merge_strategy_with_loglikelihood_function function call (train_sets[:, 2:end]()) + colloc_train_sets = [[hcat([-1], train_sets_pde[i][:, j]...) + for i in eachindex(data_colloc_loss_functions[1])] + for j in eachindex(data_colloc_loss_functions)] + + # using dataset's indvar coords as train_sets_pde and indvar coord's datafree_colloc_loss_function, create loss functions + # placeholder strategy = GridTraining(0.1), datafree_bc_loss_function and train_sets_bc must be nothing + # order of indvar coords will be same as corresponding depvar coords values in dataset provided in get_lossy() call. + pde_loss_function_points = [merge_strategy_with_loglikelihood_function( + pinnrep, + GridTraining(0.1), + data_colloc_loss_functions[i], + nothing; + train_sets_pde = colloc_train_sets[i], + train_sets_bc = nothing)[1] + for i in eachindex(data_colloc_loss_functions)] + + function L2_loss2(θ, phynewstd) + # first sum is over points losses over many equations for the same points + # second sum is over all points + pde_loglikelihoods = sum([sum([pde_loss_function(θ, + phynewstd[i]) + for (i, pde_loss_function) in enumerate(pde_loss_functions)]) + for pde_loss_functions in pde_loss_function_points]) + end + end + + # add overall functionality for BC dataset points (case of parametric BC) ? + if ((dataset_bc isa Nothing) && (dataset_pde isa Nothing)) + dataset = nothing + elseif dataset_bc isa Nothing + dataset = dataset_pde + elseif dataset_pde isa Nothing + dataset = dataset_bc + else + dataset = [vcat(dataset_pde[i], dataset_bc[i]) for i in eachindex(dataset_pde)] + end + + if discretization.param_estim && isempty(param) + throw(UndefVarError(:param)) + elseif discretization.param_estim && dataset isa Nothing + throw(UndefVarError(:dataset)) + elseif discretization.param_estim && length(l2std) != length(pinnrep.depvars) + error("L2 stds length must match number of dependant variables") + end + + # for physics loglikelihood + full_weighted_loglikelihood = pinnrep.loss_functions.full_loss_function + chain = discretization.chain + + if length(pinnrep.domains) != length(saveats) + error("Number of independent variables must match saveat inference discretization steps") + end + + # NN solutions for loglikelihood which is used for L2lossdata + Φ = pinnrep.phi + + @assert nchains≥1 "number of chains must be greater than or equal to 1" + + # remove inv params take only NN params, AHMC uses Float64 + initial_nnθ = pinnrep.flat_init_params[1:(end - length(param))] + initial_θ = collect(Float64, initial_nnθ) + + # contains only NN parameters + initial_nnθ = pinnrep.init_params + + names = ntuple(i -> pinnrep.depvars[i], length(chain)) + + #ode parameter estimation + nparameters = length(initial_θ) + ninv = length(param) + # add init_params for NN params + priors = [ + MvNormal(priorsNNw[1] * ones(nparameters), + Diagonal(abs2.(priorsNNw[2] .* ones(nparameters)))) + ] + + # append Ode params to all paramvector - initial_θ + if ninv > 0 + # shift ode params(initialise ode params by prior means) + # check if means or user specified is better + initial_θ = vcat(initial_θ, [Distributions.params(param[i])[1] for i in 1:ninv]) + priors = vcat(priors, param) + nparameters += ninv + end + + # vector in case of N-dimensional domains + strategy = discretization.strategy + + # dimensions would be total no of params,initial_nnθ for Lux namedTuples + ℓπ = PDELogTargetDensity( + nparameters, strategy, dataset, priors, [phystd, bcstd, l2std], phynewstd, + names, ninv, initial_nnθ, full_weighted_loglikelihood, newloss, Φ) + + Adaptor, Metric, targetacceptancerate = Adaptorkwargs[:Adaptor], + Adaptorkwargs[:Metric], Adaptorkwargs[:targetacceptancerate] + + # Define Hamiltonian system (nparameters ~ dimensionality of the sampling space) + metric = Metric(nparameters) + hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) + + if verbose + @printf("Current Physics Log-likelihood : %g\n", + ℓπ.full_loglikelihood(setparameters(ℓπ, initial_θ), ℓπ.allstd)) + @printf("Current Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, initial_θ)) + @printf("Current SSE against dataset Log-likelihood : %g\n", + L2LossData(ℓπ, initial_θ)) + if !(newloss isa Nothing) + @printf("Current new loss : %g\n", + ℓπ.L2_loss2(setparameters(ℓπ, initial_θ), + ℓπ.phynewstd)) + end + end + + # parallel sampling option + if nchains != 1 + + # Cache to store the chains + bpinnsols = Vector{Any}(undef, nchains) + + Threads.@threads for i in 1:nchains + # each chain has different initial NNparameter values(better posterior exploration) + initial_θ = vcat(randn(nparameters - ninv), + initial_θ[(nparameters - ninv + 1):end]) + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = integratorchoice(Integratorkwargs, initial_ϵ) + adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), + StepSizeAdaptor(targetacceptancerate, integrator)) + Kernel = AdvancedHMC.make_kernel(Kernel, integrator) + samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, adaptor; + progress = progress, verbose = verbose) + + # return a chain(basic chain),samples and stats + matrix_samples = hcat(samples...) + mcmc_chain = MCMCChains.Chains(matrix_samples') + + fullsolution = BPINNstats(mcmc_chain, samples, stats) + ensemblecurves, estimnnparams, estimated_params, timepoints = inference( + samples, pinnrep, saveat, numensemble, ℓπ) + + bpinnsols[i] = BPINNsolution( + fullsolution, ensemblecurves, estimnnparams, estimated_params, timepoints) + end + return bpinnsols + else + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = integratorchoice(Integratorkwargs, initial_ϵ) + adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), + StepSizeAdaptor(targetacceptancerate, integrator)) + + Kernel = AdvancedHMC.make_kernel(Kernel, integrator) + samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, + adaptor; progress = progress, verbose = verbose) + + # return a chain(basic chain),samples and stats + matrix_samples = hcat(samples...) + mcmc_chain = MCMCChains.Chains(matrix_samples') + + if verbose + @printf("Sampling Complete.\n") + @printf("Final Physics Log-likelihood : %g\n", + ℓπ.full_loglikelihood(setparameters(ℓπ, samples[end]), ℓπ.allstd)) + @printf("Final Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, samples[end])) + @printf("Final SSE against dataset Log-likelihood : %g\n", + L2LossData(ℓπ, samples[end])) + if !(newloss isa Nothing) + @printf("Final new loss : %g\n", + ℓπ.L2_loss2(setparameters(ℓπ, samples[end]), + ℓπ.phynewstd)) + end + end + + fullsolution = BPINNstats(mcmc_chain, samples, stats) + ensemblecurves, estimnnparams, estimated_params, timepoints = inference(samples, + pinnrep, saveats, numensemble, ℓπ) + + return BPINNsolution( + fullsolution, ensemblecurves, estimnnparams, estimated_params, timepoints) + end end diff --git a/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl b/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl index ff0c732da..a577e89b6 100644 --- a/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl +++ b/lib/BayesianNeuralPDE/src/advancedHMC_MCMC.jl @@ -103,138 +103,138 @@ Incase you are only solving the Equations for solution, do not provide dataset releases. """ function ahmc_bayesian_pinn_ode( - prob::SciMLBase.ODEProblem, chain; strategy = GridTraining, dataset = [nothing], - init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0, l2std = [0.05], - phystd = [0.05], phynewstd = [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, estim_collocate = false) - @assert !isinplace(prob) "The BPINN ODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t)." - - chain isa AbstractLuxLayer || (chain = FromFluxAdaptor()(chain)) - - strategy = strategy == GridTraining ? strategy(physdt) : strategy - - if dataset != [nothing] && - (length(dataset) < 2 || !(dataset isa Vector{<:Vector{<:AbstractFloat}})) - error("Invalid dataset. dataset would be timeseries (x̂,t) where type: Vector{Vector{AbstractFloat}") - end - - if dataset != [nothing] && param == [] - println("Dataset is only needed for Parameter Estimation + Forward Problem, not in only Forward Problem case.") - elseif dataset == [nothing] && param != [] - error("Dataset Required for Parameter Estimation.") - end - - initial_nnθ, chain, st = generate_ltd(chain, init_params) - - @assert nchains ≤ Threads.nthreads() "number of chains is greater than available threads" - @assert nchains ≥ 1 "number of chains must be greater than 1" - - # eltype(physdt) cause needs Float64 for find_good_stepsize - # Lux chain(using component array later as vector_to_parameter need namedtuple) - T = eltype(physdt) - initial_θ = getdata(ComponentArray{T}(initial_nnθ)) - - # adding ode parameter estimation - nparameters = length(initial_θ) - ninv = length(param) - priors = [ - MvNormal(T(priorsNNw[1]) * ones(T, nparameters), - Diagonal(abs2.(T(priorsNNw[2]) .* ones(T, nparameters)))), - ] - - # append Ode params to all paramvector - if ninv > 0 - # shift ode params(initialise ode params by prior means) - initial_θ = vcat(initial_θ, [Distributions.params(param[i])[1] for i in 1:ninv]) - priors = vcat(priors, param) - nparameters += ninv - end - - smodel = StatefulLuxLayer{true}(chain, nothing, st) - # dimensions would be total no of params,initial_nnθ for Lux namedTuples - ℓπ = LogTargetDensity(nparameters, prob, smodel, strategy, dataset, priors, - phystd, phynewstd, l2std, autodiff, physdt, ninv, initial_nnθ, estim_collocate) - - if verbose - @printf("Current Physics Log-likelihood: %g\n", physloglikelihood(ℓπ, initial_θ)) - @printf("Current Prior Log-likelihood: %g\n", priorweights(ℓπ, initial_θ)) - @printf("Current SSE against dataset Log-likelihood: %g\n", - L2LossData(ℓπ, initial_θ)) - if estim_collocate - @printf("Current gradient loss against dataset Log-likelihood: %g\n", - L2loss2(ℓπ, initial_θ)) - end - end - - Adaptor = Adaptorkwargs[:Adaptor] - Metric = Adaptorkwargs[:Metric] - targetacceptancerate = Adaptorkwargs[:targetacceptancerate] - - # Define Hamiltonian system (nparameters ~ dimensionality of the sampling space) - metric = Metric(nparameters) - hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) - - # parallel sampling option - if nchains != 1 - # Cache to store the chains - chains = Vector{Any}(undef, nchains) - statsc = Vector{Any}(undef, nchains) - samplesc = Vector{Any}(undef, nchains) - - Threads.@threads for i in 1:nchains - # each chain has different initial NNparameter values(better posterior exploration) - initial_θ = vcat( - randn(eltype(initial_θ), nparameters - ninv), - initial_θ[(nparameters-ninv+1):end], - ) - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = integratorchoice(Integratorkwargs, initial_ϵ) - adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), - StepSizeAdaptor(targetacceptancerate, integrator)) - - MCMC_alg = kernelchoice(Kernel, MCMCkwargs) - Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) - samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, adaptor; - progress = progress, verbose = verbose) - - samplesc[i] = samples - statsc[i] = stats - mcmc_chain = Chains(reduce(hcat, samples)') - chains[i] = mcmc_chain - end - - return chains, samplesc, statsc - else - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = integratorchoice(Integratorkwargs, initial_ϵ) - adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), - StepSizeAdaptor(targetacceptancerate, integrator)) - - MCMC_alg = kernelchoice(Kernel, MCMCkwargs) - Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) - samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, - adaptor; progress = progress, verbose = verbose) - - if verbose - println("Sampling Complete.") - @printf("Final Physics Log-likelihood: %g\n", - physloglikelihood(ℓπ, samples[end])) - @printf("Final Prior Log-likelihood: %g\n", priorweights(ℓπ, samples[end])) - @printf("Final SSE against dataset Log-likelihood: %g\n", - L2LossData(ℓπ, samples[end])) - if estim_collocate - @printf("Final gradient loss against dataset Log-likelihood: %g\n", - L2loss2(ℓπ, samples[end])) - end - end - - # return a chain(basic chain),samples and stats - matrix_samples = reshape(hcat(samples...), (length(samples[1]), length(samples), 1)) - mcmc_chain = MCMCChains.Chains(matrix_samples) - return mcmc_chain, samples, stats - end + prob::SciMLBase.ODEProblem, chain; strategy = GridTraining, dataset = [nothing], + init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0, l2std = [0.05], + phystd = [0.05], phynewstd = [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, estim_collocate = false) + @assert !isinplace(prob) "The BPINN ODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t)." + + chain isa AbstractLuxLayer || (chain = FromFluxAdaptor()(chain)) + + strategy = strategy == GridTraining ? strategy(physdt) : strategy + + if dataset != [nothing] && + (length(dataset) < 2 || !(dataset isa Vector{<:Vector{<:AbstractFloat}})) + error("Invalid dataset. dataset would be timeseries (x̂,t) where type: Vector{Vector{AbstractFloat}") + end + + if dataset != [nothing] && param == [] + println("Dataset is only needed for Parameter Estimation + Forward Problem, not in only Forward Problem case.") + elseif dataset == [nothing] && param != [] + error("Dataset Required for Parameter Estimation.") + end + + initial_nnθ, chain, st = generate_ltd(chain, init_params) + + @assert nchains≤Threads.nthreads() "number of chains is greater than available threads" + @assert nchains≥1 "number of chains must be greater than 1" + + # eltype(physdt) cause needs Float64 for find_good_stepsize + # Lux chain(using component array later as vector_to_parameter need namedtuple) + T = eltype(physdt) + initial_θ = getdata(ComponentArray{T}(initial_nnθ)) + + # adding ode parameter estimation + nparameters = length(initial_θ) + ninv = length(param) + priors = [ + MvNormal(T(priorsNNw[1]) * ones(T, nparameters), + Diagonal(abs2.(T(priorsNNw[2]) .* ones(T, nparameters)))) + ] + + # append Ode params to all paramvector + if ninv > 0 + # shift ode params(initialise ode params by prior means) + initial_θ = vcat(initial_θ, [Distributions.params(param[i])[1] for i in 1:ninv]) + priors = vcat(priors, param) + nparameters += ninv + end + + smodel = StatefulLuxLayer{true}(chain, nothing, st) + # dimensions would be total no of params,initial_nnθ for Lux namedTuples + ℓπ = LogTargetDensity(nparameters, prob, smodel, strategy, dataset, priors, + phystd, phynewstd, l2std, autodiff, physdt, ninv, initial_nnθ, estim_collocate) + + if verbose + @printf("Current Physics Log-likelihood: %g\n", physloglikelihood(ℓπ, initial_θ)) + @printf("Current Prior Log-likelihood: %g\n", priorweights(ℓπ, initial_θ)) + @printf("Current SSE against dataset Log-likelihood: %g\n", + L2LossData(ℓπ, initial_θ)) + if estim_collocate + @printf("Current gradient loss against dataset Log-likelihood: %g\n", + L2loss2(ℓπ, initial_θ)) + end + end + + Adaptor = Adaptorkwargs[:Adaptor] + Metric = Adaptorkwargs[:Metric] + targetacceptancerate = Adaptorkwargs[:targetacceptancerate] + + # Define Hamiltonian system (nparameters ~ dimensionality of the sampling space) + metric = Metric(nparameters) + hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) + + # parallel sampling option + if nchains != 1 + # Cache to store the chains + chains = Vector{Any}(undef, nchains) + statsc = Vector{Any}(undef, nchains) + samplesc = Vector{Any}(undef, nchains) + + Threads.@threads for i in 1:nchains + # each chain has different initial NNparameter values(better posterior exploration) + initial_θ = vcat( + randn(eltype(initial_θ), nparameters - ninv), + initial_θ[(nparameters - ninv + 1):end] + ) + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = integratorchoice(Integratorkwargs, initial_ϵ) + adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), + StepSizeAdaptor(targetacceptancerate, integrator)) + + MCMC_alg = kernelchoice(Kernel, MCMCkwargs) + Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) + samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, adaptor; + progress = progress, verbose = verbose) + + samplesc[i] = samples + statsc[i] = stats + mcmc_chain = Chains(reduce(hcat, samples)') + chains[i] = mcmc_chain + end + + return chains, samplesc, statsc + else + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = integratorchoice(Integratorkwargs, initial_ϵ) + adaptor = adaptorchoice(Adaptor, MassMatrixAdaptor(metric), + StepSizeAdaptor(targetacceptancerate, integrator)) + + MCMC_alg = kernelchoice(Kernel, MCMCkwargs) + Kernel = AdvancedHMC.make_kernel(MCMC_alg, integrator) + samples, stats = sample(hamiltonian, Kernel, initial_θ, draw_samples, + adaptor; progress = progress, verbose = verbose) + + if verbose + println("Sampling Complete.") + @printf("Final Physics Log-likelihood: %g\n", + physloglikelihood(ℓπ, samples[end])) + @printf("Final Prior Log-likelihood: %g\n", priorweights(ℓπ, samples[end])) + @printf("Final SSE against dataset Log-likelihood: %g\n", + L2LossData(ℓπ, samples[end])) + if estim_collocate + @printf("Final gradient loss against dataset Log-likelihood: %g\n", + L2loss2(ℓπ, samples[end])) + end + end + + # return a chain(basic chain),samples and stats + matrix_samples = reshape(hcat(samples...), (length(samples[1]), length(samples), 1)) + mcmc_chain = MCMCChains.Chains(matrix_samples) + return mcmc_chain, samples, stats + end end diff --git a/lib/BayesianNeuralPDE/src/discretize.jl b/lib/BayesianNeuralPDE/src/discretize.jl index e7fa68941..8b655a6e9 100644 --- a/lib/BayesianNeuralPDE/src/discretize.jl +++ b/lib/BayesianNeuralPDE/src/discretize.jl @@ -1,78 +1,78 @@ function get_likelihood_estimate_function(discretization::BayesianPINN) - dataset_pde, dataset_bc = discretization.dataset + dataset_pde, dataset_bc = discretization.dataset - pde_loss_functions, bc_loss_functions = merge_strategy_with_loglikelihood_function( - pinnrep, strategy, - datafree_pde_loss_functions, datafree_bc_loss_functions) + pde_loss_functions, bc_loss_functions = merge_strategy_with_loglikelihood_function( + pinnrep, strategy, + datafree_pde_loss_functions, datafree_bc_loss_functions) - # required as Physics loss also needed on the discrete dataset domain points - # data points are discrete and so by default GridTraining loss applies - # passing placeholder dx with GridTraining, it uses data points irl - datapde_loss_functions, databc_loss_functions = if dataset_bc !== nothing || - dataset_pde !== nothing - merge_strategy_with_loglikelihood_function(pinnrep, GridTraining(0.1), - datafree_pde_loss_functions, datafree_bc_loss_functions, - train_sets_pde = dataset_pde, train_sets_bc = dataset_bc) - else - nothing, nothing - end + # required as Physics loss also needed on the discrete dataset domain points + # data points are discrete and so by default GridTraining loss applies + # passing placeholder dx with GridTraining, it uses data points irl + datapde_loss_functions, databc_loss_functions = if dataset_bc !== nothing || + dataset_pde !== nothing + merge_strategy_with_loglikelihood_function(pinnrep, GridTraining(0.1), + datafree_pde_loss_functions, datafree_bc_loss_functions, + train_sets_pde = dataset_pde, train_sets_bc = dataset_bc) + else + nothing, nothing + end - # this includes losses from dataset domain points as well as discretization points - function full_loss_function(θ, allstd::Vector{Vector{Float64}}) - stdpdes, stdbcs, stdextra = allstd - # the aggregation happens on cpu even if the losses are gpu, probably fine since it's only a few of them - # SSE FOR LOSS ON GRIDPOINTS not MSE ! i, j depend on number of bcs and eqs - pde_loglikelihoods = sum([pde_loglike_function(θ, stdpdes[i]) - for (i, pde_loglike_function) in enumerate(pde_loss_functions)]) + # this includes losses from dataset domain points as well as discretization points + function full_loss_function(θ, allstd::Vector{Vector{Float64}}) + stdpdes, stdbcs, stdextra = allstd + # the aggregation happens on cpu even if the losses are gpu, probably fine since it's only a few of them + # SSE FOR LOSS ON GRIDPOINTS not MSE ! i, j depend on number of bcs and eqs + pde_loglikelihoods = sum([pde_loglike_function(θ, stdpdes[i]) + for (i, pde_loglike_function) in enumerate(pde_loss_functions)]) - bc_loglikelihoods = sum([bc_loglike_function(θ, stdbcs[j]) - for (j, bc_loglike_function) in enumerate(bc_loss_functions)]) + bc_loglikelihoods = sum([bc_loglike_function(θ, stdbcs[j]) + for (j, bc_loglike_function) in enumerate(bc_loss_functions)]) - # final newloss creation components are similar to this - if !(datapde_loss_functions isa Nothing) - pde_loglikelihoods += sum([pde_loglike_function(θ, stdpdes[j]) - for (j, pde_loglike_function) in enumerate(datapde_loss_functions)]) - end + # final newloss creation components are similar to this + if !(datapde_loss_functions isa Nothing) + pde_loglikelihoods += sum([pde_loglike_function(θ, stdpdes[j]) + for (j, pde_loglike_function) in enumerate(datapde_loss_functions)]) + end - if !(databc_loss_functions isa Nothing) - bc_loglikelihoods += sum([bc_loglike_function(θ, stdbcs[j]) - for (j, bc_loglike_function) in enumerate(databc_loss_functions)]) - end + if !(databc_loss_functions isa Nothing) + bc_loglikelihoods += sum([bc_loglike_function(θ, stdbcs[j]) + for (j, bc_loglike_function) in enumerate(databc_loss_functions)]) + end - # this is kind of a hack, and means that whenever the outer function is evaluated the increment goes up, even if it's not being optimized - # that's why we prefer the user to maintain the increment in the outer loop callback during optimization - @ignore_derivatives if self_increment - iteration[] += 1 - end + # this is kind of a hack, and means that whenever the outer function is evaluated the increment goes up, even if it's not being optimized + # that's why we prefer the user to maintain the increment in the outer loop callback during optimization + @ignore_derivatives if self_increment + iteration[] += 1 + end - @ignore_derivatives begin - reweight_losses_func(θ, pde_loglikelihoods, - bc_loglikelihoods) - end + @ignore_derivatives begin + reweight_losses_func(θ, pde_loglikelihoods, + bc_loglikelihoods) + end - weighted_pde_loglikelihood = adaloss.pde_loss_weights .* pde_loglikelihoods - weighted_bc_loglikelihood = adaloss.bc_loss_weights .* bc_loglikelihoods + weighted_pde_loglikelihood = adaloss.pde_loss_weights .* pde_loglikelihoods + weighted_bc_loglikelihood = adaloss.bc_loss_weights .* bc_loglikelihoods - sum_weighted_pde_loglikelihood = sum(weighted_pde_loglikelihood) - sum_weighted_bc_loglikelihood = sum(weighted_bc_loglikelihood) - weighted_loglikelihood_before_additional = sum_weighted_pde_loglikelihood + - sum_weighted_bc_loglikelihood + sum_weighted_pde_loglikelihood = sum(weighted_pde_loglikelihood) + sum_weighted_bc_loglikelihood = sum(weighted_bc_loglikelihood) + weighted_loglikelihood_before_additional = sum_weighted_pde_loglikelihood + + sum_weighted_bc_loglikelihood - full_weighted_loglikelihood = if additional_loss isa Nothing - weighted_loglikelihood_before_additional - else - (θ_, p_) = param_estim ? (θ.depvar, θ.p) : (θ, nothing) - _additional_loss = additional_loss(phi, θ_, p_) - _additional_loglikelihood = logpdf(Normal(0, stdextra), _additional_loss) + full_weighted_loglikelihood = if additional_loss isa Nothing + weighted_loglikelihood_before_additional + else + (θ_, p_) = param_estim ? (θ.depvar, θ.p) : (θ, nothing) + _additional_loss = additional_loss(phi, θ_, p_) + _additional_loglikelihood = logpdf(Normal(0, stdextra), _additional_loss) - weighted_additional_loglikelihood = adaloss.additional_loss_weights[1] * - _additional_loglikelihood + weighted_additional_loglikelihood = adaloss.additional_loss_weights[1] * + _additional_loglikelihood - weighted_loglikelihood_before_additional + weighted_additional_loglikelihood - end + weighted_loglikelihood_before_additional + weighted_additional_loglikelihood + end - return full_weighted_loglikelihood - end + return full_weighted_loglikelihood + end - return full_loss_function + return full_loss_function end diff --git a/lib/BayesianNeuralPDE/src/pinn_types.jl b/lib/BayesianNeuralPDE/src/pinn_types.jl index c71194fd7..bc0fb5fc5 100644 --- a/lib/BayesianNeuralPDE/src/pinn_types.jl +++ b/lib/BayesianNeuralPDE/src/pinn_types.jl @@ -17,17 +17,17 @@ the ones mentioned below. inverse problem solving. """ @concrete struct BayesianPINN <: AbstractPINN - pinn <: PhysicsInformedNN - dataset::Any + pinn <: PhysicsInformedNN + dataset::Any end function Base.getproperty(pinn::BayesianPINN, name::Symbol) - name === :dataset && return getfield(pinn, :dataset) - name === :pinn && return getfield(pinn, :pinn) - return getproperty(pinn.pinn, name) + name === :dataset && return getfield(pinn, :dataset) + name === :pinn && return getfield(pinn, :pinn) + return getproperty(pinn.pinn, name) end function BayesianPINN(args...; dataset = nothing, kwargs...) - dataset === nothing && (dataset = (nothing, nothing)) - return BayesianPINN(PhysicsInformedNN(args...; kwargs...), dataset) + dataset === nothing && (dataset = (nothing, nothing)) + return BayesianPINN(PhysicsInformedNN(args...; kwargs...), dataset) end diff --git a/lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl b/lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl index 296e4f216..d8456e051 100644 --- a/lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl +++ b/lib/BayesianNeuralPDE/test/BPINN_PDE_tests.jl @@ -1,529 +1,529 @@ -@testitem "BPINN PDE I: 1D Periodic System" tags = [:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum - - Random.seed!(100) - - @parameters t - @variables u(..) - Dt = Differential(t) - eq = Dt(u(t)) - cospi(2t) ~ 0 - bcs = [u(0.0) ~ 0.0] - domains = [t ∈ Interval(0.0, 2.0)] - - chainl = Chain(Dense(1, 6, tanh), Dense(6, 1)) - initl, st = Lux.setup(Random.default_rng(), chainl) - @named pde_system = PDESystem(eq, bcs, domains, [t], [u(t)]) - - # non adaptive case - discretization = BayesianPINN([chainl], GridTraining([0.01])) - - sol1 = ahmc_bayesian_pinn_pde( - pde_system, discretization; draw_samples = 1500, bcstd = [0.01], - phystd = [0.01], priorsNNw = (0.0, 1.0), saveats = [1 / 50.0]) - - analytic_sol_func(u0, t) = u0 + sinpi(2t) / (2pi) - ts = vec(sol1.timepoints[1]) - u_real = [analytic_sol_func(0.0, t) for t in ts] - u_predict = pmean(sol1.ensemblesol[1]) - - # absol tests - @test mean(abs, u_predict .- u_real) < 5e-2 +@testitem "BPINN PDE I: 1D Periodic System" tags=[:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum + + Random.seed!(100) + + @parameters t + @variables u(..) + Dt = Differential(t) + eq = Dt(u(t)) - cospi(2t) ~ 0 + bcs = [u(0.0) ~ 0.0] + domains = [t ∈ Interval(0.0, 2.0)] + + chainl = Chain(Dense(1, 6, tanh), Dense(6, 1)) + initl, st = Lux.setup(Random.default_rng(), chainl) + @named pde_system = PDESystem(eq, bcs, domains, [t], [u(t)]) + + # non adaptive case + discretization = BayesianPINN([chainl], GridTraining([0.01])) + + sol1 = ahmc_bayesian_pinn_pde( + pde_system, discretization; draw_samples = 1500, bcstd = [0.01], + phystd = [0.01], priorsNNw = (0.0, 1.0), saveats = [1 / 50.0]) + + analytic_sol_func(u0, t) = u0 + sinpi(2t) / (2pi) + ts = vec(sol1.timepoints[1]) + u_real = [analytic_sol_func(0.0, t) for t in ts] + u_predict = pmean(sol1.ensemblesol[1]) + + # absol tests + @test mean(abs, u_predict .- u_real) < 5e-2 end -@testitem "BPINN PDE II: 1D ODE" tags = [:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum +@testitem "BPINN PDE II: 1D ODE" tags=[:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum - Random.seed!(100) + Random.seed!(100) - @parameters θ - @variables u(..) - Dθ = Differential(θ) + @parameters θ + @variables u(..) + Dθ = Differential(θ) - # 1D ODE - eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) - - u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3))) + # 1D ODE + eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) - + u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3))) - # Initial and boundary conditions - bcs = [u(0.0) ~ 1.0f0] + # Initial and boundary conditions + bcs = [u(0.0) ~ 1.0f0] - # Space and time domains - domains = [θ ∈ Interval(0.0f0, 1.0f0)] + # Space and time domains + domains = [θ ∈ Interval(0.0f0, 1.0f0)] - # Discretization - dt = 0.1f0 + # Discretization + dt = 0.1f0 - # Neural network - chain = Chain(Dense(1, 12, σ), Dense(12, 1)) + # Neural network + chain = Chain(Dense(1, 12, σ), Dense(12, 1)) - discretization = BayesianPINN([chain], GridTraining([0.01])) + discretization = BayesianPINN([chain], GridTraining([0.01])) - @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) + @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) - sol1 = ahmc_bayesian_pinn_pde( - pde_system, discretization; draw_samples = 500, bcstd = [0.1], - phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) + sol1 = ahmc_bayesian_pinn_pde( + pde_system, discretization; draw_samples = 500, bcstd = [0.1], + phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) - analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2 - ts = sol1.timepoints[1] - u_real = vec([analytic_sol_func(t) for t in ts]) - u_predict = pmean(sol1.ensemblesol[1]) - @test u_predict ≈ u_real atol = 0.8 + analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2 + ts = sol1.timepoints[1] + u_real = vec([analytic_sol_func(t) for t in ts]) + u_predict = pmean(sol1.ensemblesol[1]) + @test u_predict≈u_real atol=0.8 end -@testitem "BPINN PDE III: 3rd Degree ODE" tags = [:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum +@testitem "BPINN PDE III: 3rd Degree ODE" tags=[:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum - Random.seed!(100) + Random.seed!(100) - @parameters x - @variables u(..), Dxu(..), Dxxu(..), O1(..), O2(..) - Dxxx = Differential(x)^3 - Dx = Differential(x) + @parameters x + @variables u(..), Dxu(..), Dxxu(..), O1(..), O2(..) + Dxxx = Differential(x)^3 + Dx = Differential(x) - # ODE - eq = Dx(Dxxu(x)) ~ cospi(x) + # ODE + eq = Dx(Dxxu(x)) ~ cospi(x) - # Initial and boundary conditions - ep = (cbrt(eps(eltype(Float64))))^2 / 6 + # Initial and boundary conditions + ep = (cbrt(eps(eltype(Float64))))^2 / 6 - bcs = [ - u(0.0) ~ 0.0, - u(1.0) ~ cospi(1.0), - Dxu(1.0) ~ 1.0, - Dxu(x) ~ Dx(u(x)) + ep * O1(x), - Dxxu(x) ~ Dx(Dxu(x)) + ep * O2(x), - ] + bcs = [ + u(0.0) ~ 0.0, + u(1.0) ~ cospi(1.0), + Dxu(1.0) ~ 1.0, + Dxu(x) ~ Dx(u(x)) + ep * O1(x), + Dxxu(x) ~ Dx(Dxu(x)) + ep * O2(x) + ] - # Space and time domains - domains = [x ∈ Interval(0.0, 1.0)] + # Space and time domains + domains = [x ∈ Interval(0.0, 1.0)] - # Neural network - chain = [ - Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), - Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), - Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), - Chain(Dense(1, 4, tanh), Dense(4, 1)), - Chain(Dense(1, 4, tanh), Dense(4, 1)), - ] + # Neural network + chain = [ + Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), + Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), + Chain(Dense(1, 10, tanh), Dense(10, 10, tanh), Dense(10, 1)), + Chain(Dense(1, 4, tanh), Dense(4, 1)), + Chain(Dense(1, 4, tanh), Dense(4, 1)) + ] - discretization = BayesianPINN(chain, GridTraining(0.01)) + discretization = BayesianPINN(chain, GridTraining(0.01)) - @named pde_system = PDESystem(eq, bcs, domains, [x], - [u(x), Dxu(x), Dxxu(x), O1(x), O2(x)]) + @named pde_system = PDESystem(eq, bcs, domains, [x], + [u(x), Dxu(x), Dxxu(x), O1(x), O2(x)]) - sol1 = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200, - bcstd = [0.01, 0.01, 0.01, 0.01, 0.01], phystd = [0.005], - priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) + sol1 = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200, + bcstd = [0.01, 0.01, 0.01, 0.01, 0.01], phystd = [0.005], + priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) - analytic_sol_func(x) = (π * x * (-x + (π^2) * (2 * x - 3) + 1) - sinpi(x)) / (π^3) + analytic_sol_func(x) = (π * x * (-x + (π^2) * (2 * x - 3) + 1) - sinpi(x)) / (π^3) - u_predict = pmean(sol1.ensemblesol[1]) - xs = vec(sol1.timepoints[1]) - u_real = [analytic_sol_func(x) for x in xs] - @test u_predict ≈ u_real atol = 0.5 + u_predict = pmean(sol1.ensemblesol[1]) + xs = vec(sol1.timepoints[1]) + u_real = [analytic_sol_func(x) for x in xs] + @test u_predict≈u_real atol=0.5 end -@testitem "BPINN PDE IV: 2D Poisson" tags = [:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum +@testitem "BPINN PDE IV: 2D Poisson" tags=[:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum - Random.seed!(100) + Random.seed!(100) - @parameters x y - @variables u(..) - Dxx = Differential(x)^2 - Dyy = Differential(y)^2 + @parameters x y + @variables u(..) + Dxx = Differential(x)^2 + Dyy = Differential(y)^2 - # 2D PDE - eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y) + # 2D PDE + eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y) - # Boundary conditions - bcs = [ - u(0, y) ~ 0.0, - u(1, y) ~ 0.0, - u(x, 0) ~ 0.0, - u(x, 1) ~ 0.0, - ] + # Boundary conditions + bcs = [ + u(0, y) ~ 0.0, + u(1, y) ~ 0.0, + u(x, 0) ~ 0.0, + u(x, 1) ~ 0.0 + ] - # Space and time domains - domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)] + # Space and time domains + domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)] - # Discretization - dt = 0.1f0 + # Discretization + dt = 0.1f0 - # Neural network - chain = Chain(Dense(2, 9, σ), Dense(9, 9, σ), Dense(9, 1)) + # Neural network + chain = Chain(Dense(2, 9, σ), Dense(9, 9, σ), Dense(9, 1)) - dx = 0.04 - discretization = BayesianPINN([chain], GridTraining(dx)) + dx = 0.04 + discretization = BayesianPINN([chain], GridTraining(dx)) - @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) + @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)]) - sol = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200, - bcstd = [0.003, 0.003, 0.003, 0.003], phystd = [0.003], - priorsNNw = (0.0, 10.0), saveats = [1 / 100.0, 1 / 100.0]) + sol = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 200, + bcstd = [0.003, 0.003, 0.003, 0.003], phystd = [0.003], + priorsNNw = (0.0, 10.0), saveats = [1 / 100.0, 1 / 100.0]) - xs = sol.timepoints[1] - analytic_sol_func(x, y) = (sinpi(x) * sinpi(y)) / (2pi^2) + xs = sol.timepoints[1] + analytic_sol_func(x, y) = (sinpi(x) * sinpi(y)) / (2pi^2) - u_predict = pmean(sol.ensemblesol[1]) - u_real = [analytic_sol_func(xs[:, i][1], xs[:, i][2]) for i in 1:length(xs[1, :])] - @test u_predict ≈ u_real rtol = 0.5 + u_predict = pmean(sol.ensemblesol[1]) + u_real = [analytic_sol_func(xs[:, i][1], xs[:, i][2]) for i in 1:length(xs[1, :])] + @test u_predict≈u_real rtol=0.5 end -@testitem "BPINN PDE: Translating from Flux" tags = [:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum - using Flux: Flux +@testitem "BPINN PDE: Translating from Flux" tags=[:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum + using Flux: Flux - Random.seed!(100) + Random.seed!(100) - @parameters θ - @variables u(..) - Dθ = Differential(θ) + @parameters θ + @variables u(..) + Dθ = Differential(θ) - # 1D ODE - eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) - - u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3))) + # 1D ODE + eq = Dθ(u(θ)) ~ θ^3 + 2.0f0 * θ + (θ^2) * ((1.0f0 + 3 * (θ^2)) / (1.0f0 + θ + (θ^3))) - + u(θ) * (θ + ((1.0f0 + 3.0f0 * (θ^2)) / (1.0f0 + θ + θ^3))) - # Initial and boundary conditions - bcs = [u(0.0) ~ 1.0f0] + # Initial and boundary conditions + bcs = [u(0.0) ~ 1.0f0] - # Space and time domains - domains = [θ ∈ Interval(0.0f0, 1.0f0)] + # Space and time domains + domains = [θ ∈ Interval(0.0f0, 1.0f0)] - # Neural network - chain = Flux.Chain(Flux.Dense(1, 12, Flux.σ), Flux.Dense(12, 1)) + # Neural network + chain = Flux.Chain(Flux.Dense(1, 12, Flux.σ), Flux.Dense(12, 1)) - discretization = BayesianPINN([chain], GridTraining([0.01])) - @test discretization.chain[1] isa Lux.AbstractLuxLayer + discretization = BayesianPINN([chain], GridTraining([0.01])) + @test discretization.chain[1] isa Lux.AbstractLuxLayer - @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) + @named pde_system = PDESystem(eq, bcs, domains, [θ], [u]) - sol = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 500, - bcstd = [0.1], phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) + sol = ahmc_bayesian_pinn_pde(pde_system, discretization; draw_samples = 500, + bcstd = [0.1], phystd = [0.05], priorsNNw = (0.0, 10.0), saveats = [1 / 100.0]) - analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2 - ts = sol.timepoints[1] - u_real = vec([analytic_sol_func(t) for t in ts]) - u_predict = pmean(sol.ensemblesol[1]) + analytic_sol_func(t) = exp(-(t^2) / 2) / (1 + t + t^3) + t^2 + ts = sol.timepoints[1] + u_real = vec([analytic_sol_func(t) for t in ts]) + u_predict = pmean(sol.ensemblesol[1]) - @test u_predict ≈ u_real atol = 0.8 + @test u_predict≈u_real atol=0.8 end -@testitem "BPINN PDE Inv I: 1D Periodic System" tags = [:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum - - Random.seed!(100) - - @parameters t p - @variables u(..) - - Dt = Differential(t) - eqs = Dt(u(t)) - cos(p * t) ~ 0 - bcs = [u(0) ~ 0.0] - domains = [t ∈ Interval(0.0, 2.0)] - - chainl = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) - initl, st = Lux.setup(Random.default_rng(), chainl) - - @named pde_system = PDESystem(eqs, - bcs, - domains, - [t], - [u(t)], - [p], - defaults = Dict([p => 4.0])) - - analytic_sol_func1(u0, t) = u0 + sinpi(2t) / (2π) - timepoints = collect(0.0:(1/100.0):2.0) - u = [analytic_sol_func1(0.0, timepoint) for timepoint in timepoints] - u = u .+ (u .* 0.2) .* randn(size(u)) - dataset = [hcat(u, timepoints)] - - # BPINNs are formulated with a mesh that must stay the same throughout sampling (as of now) - @testset "$(nameof(typeof(strategy)))" for strategy in [ - # StochasticTraining(200), - # QuasiRandomTraining(200), - GridTraining([0.02]), - ] - discretization = BayesianPINN([chainl], strategy; param_estim = true, - dataset = [dataset, nothing]) - - sol1 = ahmc_bayesian_pinn_pde(pde_system, - discretization; - draw_samples = 1500, - bcstd = [0.02], - phystd = [0.02], l2std = [0.02], - priorsNNw = (0.0, 1.0), - saveats = [1 / 50.0], - param = [LogNormal(6.0, 0.5)]) - - param = 2 * π - ts = vec(sol1.timepoints[1]) - u_real = [analytic_sol_func1(0.0, t) for t in ts] - u_predict = pmean(sol1.ensemblesol[1]) - - @test mean(abs, u_predict .- u_real) < 5e-2 - @test sol1.estimated_de_params[1] ≈ param rtol = 0.1 - end +@testitem "BPINN PDE Inv I: 1D Periodic System" tags=[:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum + + Random.seed!(100) + + @parameters t p + @variables u(..) + + Dt = Differential(t) + eqs = Dt(u(t)) - cos(p * t) ~ 0 + bcs = [u(0) ~ 0.0] + domains = [t ∈ Interval(0.0, 2.0)] + + chainl = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) + initl, st = Lux.setup(Random.default_rng(), chainl) + + @named pde_system = PDESystem(eqs, + bcs, + domains, + [t], + [u(t)], + [p], + defaults = Dict([p => 4.0])) + + analytic_sol_func1(u0, t) = u0 + sinpi(2t) / (2π) + timepoints = collect(0.0:(1 / 100.0):2.0) + u = [analytic_sol_func1(0.0, timepoint) for timepoint in timepoints] + u = u .+ (u .* 0.2) .* randn(size(u)) + dataset = [hcat(u, timepoints)] + + # BPINNs are formulated with a mesh that must stay the same throughout sampling (as of now) + @testset "$(nameof(typeof(strategy)))" for strategy in [ + # StochasticTraining(200), + # QuasiRandomTraining(200), + GridTraining([0.02]) + ] + discretization = BayesianPINN([chainl], strategy; param_estim = true, + dataset = [dataset, nothing]) + + sol1 = ahmc_bayesian_pinn_pde(pde_system, + discretization; + draw_samples = 1500, + bcstd = [0.02], + phystd = [0.02], l2std = [0.02], + priorsNNw = (0.0, 1.0), + saveats = [1 / 50.0], + param = [LogNormal(6.0, 0.5)]) + + param = 2 * π + ts = vec(sol1.timepoints[1]) + u_real = [analytic_sol_func1(0.0, t) for t in ts] + u_predict = pmean(sol1.ensemblesol[1]) + + @test mean(abs, u_predict .- u_real) < 5e-2 + @test sol1.estimated_de_params[1]≈param rtol=0.1 + end end -@testitem "BPINN PDE Inv II: Lorenz System" tags = [:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum - - Random.seed!(100) - - @parameters t, σ_ - @variables x(..), y(..), z(..) - Dt = Differential(t) - eqs = [ - Dt(x(t)) ~ σ_ * (y(t) - x(t)), - Dt(y(t)) ~ x(t) * (28.0 - z(t)) - y(t), - Dt(z(t)) ~ x(t) * y(t) - 8.0 / 3.0 * z(t), - ] - - bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0] - domains = [t ∈ Interval(0.0, 1.0)] - - input_ = length(domains) - n = 7 - chain = [ - Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)), - Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)), - Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)), - ] - - # Generate Data - function lorenz!(du, u, p, t) - du[1] = 10.0 * (u[2] - u[1]) - du[2] = u[1] * (28.0 - u[3]) - u[2] - du[3] = u[1] * u[2] - (8.0 / 3.0) * u[3] - end - - u0 = [1.0; 0.0; 0.0] - tspan = (0.0, 1.0) - prob = ODEProblem(lorenz!, u0, tspan) - sol = solve(prob, Tsit5(), dt = 0.01, saveat = 0.05) - ts = sol.t - us = hcat(sol.u...) - us = us .+ ((0.05 .* randn(size(us))) .* us) - ts_ = hcat(ts...)[1, :] - dataset = [hcat(us[i, :], ts_) for i in 1:3] - - discretization = BayesianPINN(chain, GridTraining([0.01]); param_estim = true, - dataset = [dataset, nothing]) - - @named pde_system = PDESystem(eqs, bcs, domains, - [t], [x(t), y(t), z(t)], [σ_], defaults = Dict([p => 1.0 for p in [σ_]])) - - sol1 = ahmc_bayesian_pinn_pde(pde_system, - discretization; - draw_samples = 50, - bcstd = [0.3, 0.3, 0.3], - phystd = [0.1, 0.1, 0.1], - l2std = [1, 1, 1], - priorsNNw = (0.0, 1.0), - saveats = [0.01], - param = [Normal(12.0, 2)]) - - idealp = 10.0 - p_ = sol1.estimated_de_params[1] - @test sum(abs, pmean(p_) - 10.00) < 0.3 * idealp[1] - # @test sum(abs, pmean(p_[2]) - (8 / 3)) < 0.3 * idealp[2] +@testitem "BPINN PDE Inv II: Lorenz System" tags=[:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum + + Random.seed!(100) + + @parameters t, σ_ + @variables x(..), y(..), z(..) + Dt = Differential(t) + eqs = [ + Dt(x(t)) ~ σ_ * (y(t) - x(t)), + Dt(y(t)) ~ x(t) * (28.0 - z(t)) - y(t), + Dt(z(t)) ~ x(t) * y(t) - 8.0 / 3.0 * z(t) + ] + + bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0] + domains = [t ∈ Interval(0.0, 1.0)] + + input_ = length(domains) + n = 7 + chain = [ + Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)), + Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)), + Chain(Dense(input_, n, tanh), Dense(n, n, tanh), Dense(n, 1)) + ] + + # Generate Data + function lorenz!(du, u, p, t) + du[1] = 10.0 * (u[2] - u[1]) + du[2] = u[1] * (28.0 - u[3]) - u[2] + du[3] = u[1] * u[2] - (8.0 / 3.0) * u[3] + end + + u0 = [1.0; 0.0; 0.0] + tspan = (0.0, 1.0) + prob = ODEProblem(lorenz!, u0, tspan) + sol = solve(prob, Tsit5(), dt = 0.01, saveat = 0.05) + ts = sol.t + us = hcat(sol.u...) + us = us .+ ((0.05 .* randn(size(us))) .* us) + ts_ = hcat(ts...)[1, :] + dataset = [hcat(us[i, :], ts_) for i in 1:3] + + discretization = BayesianPINN(chain, GridTraining([0.01]); param_estim = true, + dataset = [dataset, nothing]) + + @named pde_system = PDESystem(eqs, bcs, domains, + [t], [x(t), y(t), z(t)], [σ_], defaults = Dict([p => 1.0 for p in [σ_]])) + + sol1 = ahmc_bayesian_pinn_pde(pde_system, + discretization; + draw_samples = 50, + bcstd = [0.3, 0.3, 0.3], + phystd = [0.1, 0.1, 0.1], + l2std = [1, 1, 1], + priorsNNw = (0.0, 1.0), + saveats = [0.01], + param = [Normal(12.0, 2)]) + + idealp = 10.0 + p_ = sol1.estimated_de_params[1] + @test sum(abs, pmean(p_) - 10.00) < 0.3 * idealp[1] + # @test sum(abs, pmean(p_[2]) - (8 / 3)) < 0.3 * idealp[2] end -@testitem "BPINN PDE Inv III: Improved Parametric Kuromo-Sivashinsky Equation solve" tags = [:pdebpinn] begin - using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, - AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, - ComponentArrays - import ModelingToolkit: Interval, infimum, supremum - - Random.seed!(100) - - function recur_expression(exp, Dict_differentials) - for in_exp in exp.args - if !(in_exp isa Expr) - # skip +,== symbols, characters etc - continue - - elseif in_exp.args[1] isa ModelingToolkit.Differential - # first symbol of differential term - # Dict_differentials for masking differential terms - # and resubstituting differentials in equations after putting in interpolations - # temp = in_exp.args[end] - Dict_differentials[eval(in_exp)] = Symbolics.variable("diff_$(length(Dict_differentials) + 1)") - return - else - recur_expression(in_exp, Dict_differentials) - end - end - end - - @parameters x, t, α - @variables u(..) - Dt = Differential(t) - Dx = Differential(x) - Dx2 = Differential(x)^2 - Dx3 = Differential(x)^3 - Dx4 = Differential(x)^4 - - # α = 1 (KS equation to be parametric in a) - β = 4 - γ = 1 - eq = Dt(u(x, t)) + u(x, t) * Dx(u(x, t)) + α * Dx2(u(x, t)) + β * Dx3(u(x, t)) + γ * Dx4(u(x, t)) ~ 0 - - u_analytic(x, t; z = -x / 2 + t) = 11 + 15 * tanh(z) - 15 * tanh(z)^2 - 15 * tanh(z)^3 - du(x, t; z = -x / 2 + t) = 15 / 2 * (tanh(z) + 1) * (3 * tanh(z) - 1) * sech(z)^2 - - bcs = [u(x, 0) ~ u_analytic(x, 0), - u(-10, t) ~ u_analytic(-10, t), - u(10, t) ~ u_analytic(10, t), - Dx(u(-10, t)) ~ du(-10, t), - Dx(u(10, t)) ~ du(10, t)] - - # Space and time domains - domains = [x ∈ Interval(-10.0, 10.0), - t ∈ Interval(0.0, 1.0)] - - # Discretization - dx = 0.4 - dt = 0.2 - - # Function to compute analytical solution at a specific point (x, t) - function u_analytic_point(x, t) - z = -x / 2 + t - return 11 + 15 * tanh(z) - 15 * tanh(z)^2 - 15 * tanh(z)^3 - end - - # Function to generate the dataset matrix - function generate_dataset_matrix(domains, dx, dt, xlim, tlim) - x_values = xlim[1]:dx:xlim[2] - t_values = tlim[1]:dt:tlim[2] - - dataset = [] - - for t in t_values - for x in x_values - u_value = u_analytic_point(x, t) - push!(dataset, [u_value, x, t]) - end - end - - return vcat([data' for data in dataset]...) - end - - # considering sparse dataset from half of x's domain - datasetpde_new = [generate_dataset_matrix(domains, dx, dt, [-10, 0], [0.0, 1.0])] - - # Adding Gaussian noise with a 0.8 std - noisydataset_new = deepcopy(datasetpde_new) - noisydataset_new[1][:, 1] = noisydataset_new[1][:, 1] .+ - (randn(size(noisydataset_new[1][:, 1])) .* 0.8) - - # Neural network - chain = Lux.Chain(Lux.Dense(2, 8, Lux.tanh), - Lux.Dense(8, 8, Lux.tanh), - Lux.Dense(8, 1)) - - # Discretization for old and new models - discretization = NeuralPDE.BayesianPINN([chain], - GridTraining([dx, dt]), param_estim = true, dataset = [noisydataset_new, nothing]) - - # let α default to 2.0 - @named pde_system = PDESystem(eq, - bcs, - domains, - [x, t], - [u(x, t)], - [α], - defaults = Dict([α => 2.0])) - - # neccesarry for loss function construction (involves Operator masking) - eqs = pde_system.eqs - Dict_differentials = Dict() - exps = toexpr.(eqs) - nullobj = [recur_expression(exp, Dict_differentials) for exp in exps] - - # Dict_differentials is now ; - # Dict{Any, Any} with 5 entries: - # Differential(x)(Differential(x)(u(x, t))) => diff_5 - # Differential(x)(Differential(x)(Differential(x)(u(x… => diff_1 - # Differential(x)(Differential(x)(Differential(x)(Dif… => diff_2 - # Differential(x)(u(x, t)) => diff_4 - # Differential(t)(u(x, t)) => diff_3 - - # using HMC algorithm due to convergence, stability, time of training. (refer to mcmc chain plots) - # choice of std for objectives is very important - # pass in Dict_differentials, phystdnew arguments when using the new model - - sol_new = ahmc_bayesian_pinn_pde(pde_system, - discretization; - draw_samples = 150, - bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], phynewstd = [0.4], - phystd = [0.2], l2std = [0.8], param = [Distributions.Normal(2.0, 2)], - priorsNNw = (0.0, 1.0), - saveats = [1 / 100.0, 1 / 100.0], - Dict_differentials = Dict_differentials) - - sol_old = ahmc_bayesian_pinn_pde(pde_system, - discretization; - draw_samples = 150, - bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], - phystd = [0.2], l2std = [0.8], param = [Distributions.Normal(2.0, 2)], - priorsNNw = (0.0, 1.0), - saveats = [1 / 100.0, 1 / 100.0]) - - phi = discretization.phi[1] - xs, ts = [infimum(d.domain):dx:supremum(d.domain) - for (d, dx) in zip(domains, [dx / 10, dt])] - u_real = [[u_analytic(x, t) for x in xs] for t in ts] - - u_predict_new = [[first(pmean(phi([x, t], sol_new.estimated_nn_params[1]))) for x in xs] - for t in ts] - - diff_u_new = [[abs(u_analytic(x, t) - - first(pmean(phi([x, t], sol_new.estimated_nn_params[1])))) - for x in xs] - for t in ts] - - u_predict_old = [[first(pmean(phi([x, t], sol_old.estimated_nn_params[1]))) for x in xs] - for t in ts] - diff_u_old = [[abs(u_analytic(x, t) - - first(pmean(phi([x, t], sol_old.estimated_nn_params[1])))) - for x in xs] - for t in ts] - - unsafe_comparisons(true) - @test all(all, [((diff_u_new[i]) .^ 2 .< 0.8) for i in 1:6]) == true - @test all(all, [((diff_u_old[i]) .^ 2 .< 0.8) for i in 1:6]) == false - - MSE_new = [mean(abs2, diff_u_new[i]) for i in 1:6] - MSE_old = [mean(abs2, diff_u_old[i]) for i in 1:6] - @test (MSE_new .< MSE_old) == [1, 1, 1, 1, 1, 1] - - param_new = sol_new.estimated_de_params[1] - param_old = sol_old.estimated_de_params[1] - α = 1 - @test abs(param_new - α) < 0.2 * α - @test abs(param_new - α) < abs(param_old - α) +@testitem "BPINN PDE Inv III: Improved Parametric Kuromo-Sivashinsky Equation solve" tags=[:pdebpinn] begin + using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq, + AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements, + ComponentArrays + import ModelingToolkit: Interval, infimum, supremum + + Random.seed!(100) + + function recur_expression(exp, Dict_differentials) + for in_exp in exp.args + if !(in_exp isa Expr) + # skip +,== symbols, characters etc + continue + + elseif in_exp.args[1] isa ModelingToolkit.Differential + # first symbol of differential term + # Dict_differentials for masking differential terms + # and resubstituting differentials in equations after putting in interpolations + # temp = in_exp.args[end] + Dict_differentials[eval(in_exp)] = Symbolics.variable("diff_$(length(Dict_differentials) + 1)") + return + else + recur_expression(in_exp, Dict_differentials) + end + end + end + + @parameters x, t, α + @variables u(..) + Dt = Differential(t) + Dx = Differential(x) + Dx2 = Differential(x)^2 + Dx3 = Differential(x)^3 + Dx4 = Differential(x)^4 + + # α = 1 (KS equation to be parametric in a) + β = 4 + γ = 1 + eq = Dt(u(x, t)) + u(x, t) * Dx(u(x, t)) + α * Dx2(u(x, t)) + β * Dx3(u(x, t)) + γ * Dx4(u(x, t)) ~ 0 + + u_analytic(x, t; z = -x / 2 + t) = 11 + 15 * tanh(z) - 15 * tanh(z)^2 - 15 * tanh(z)^3 + du(x, t; z = -x / 2 + t) = 15 / 2 * (tanh(z) + 1) * (3 * tanh(z) - 1) * sech(z)^2 + + bcs = [u(x, 0) ~ u_analytic(x, 0), + u(-10, t) ~ u_analytic(-10, t), + u(10, t) ~ u_analytic(10, t), + Dx(u(-10, t)) ~ du(-10, t), + Dx(u(10, t)) ~ du(10, t)] + + # Space and time domains + domains = [x ∈ Interval(-10.0, 10.0), + t ∈ Interval(0.0, 1.0)] + + # Discretization + dx = 0.4 + dt = 0.2 + + # Function to compute analytical solution at a specific point (x, t) + function u_analytic_point(x, t) + z = -x / 2 + t + return 11 + 15 * tanh(z) - 15 * tanh(z)^2 - 15 * tanh(z)^3 + end + + # Function to generate the dataset matrix + function generate_dataset_matrix(domains, dx, dt, xlim, tlim) + x_values = xlim[1]:dx:xlim[2] + t_values = tlim[1]:dt:tlim[2] + + dataset = [] + + for t in t_values + for x in x_values + u_value = u_analytic_point(x, t) + push!(dataset, [u_value, x, t]) + end + end + + return vcat([data' for data in dataset]...) + end + + # considering sparse dataset from half of x's domain + datasetpde_new = [generate_dataset_matrix(domains, dx, dt, [-10, 0], [0.0, 1.0])] + + # Adding Gaussian noise with a 0.8 std + noisydataset_new = deepcopy(datasetpde_new) + noisydataset_new[1][:, 1] = noisydataset_new[1][:, 1] .+ + (randn(size(noisydataset_new[1][:, 1])) .* 0.8) + + # Neural network + chain = Lux.Chain(Lux.Dense(2, 8, Lux.tanh), + Lux.Dense(8, 8, Lux.tanh), + Lux.Dense(8, 1)) + + # Discretization for old and new models + discretization = NeuralPDE.BayesianPINN([chain], + GridTraining([dx, dt]), param_estim = true, dataset = [noisydataset_new, nothing]) + + # let α default to 2.0 + @named pde_system = PDESystem(eq, + bcs, + domains, + [x, t], + [u(x, t)], + [α], + defaults = Dict([α => 2.0])) + + # neccesarry for loss function construction (involves Operator masking) + eqs = pde_system.eqs + Dict_differentials = Dict() + exps = toexpr.(eqs) + nullobj = [recur_expression(exp, Dict_differentials) for exp in exps] + + # Dict_differentials is now ; + # Dict{Any, Any} with 5 entries: + # Differential(x)(Differential(x)(u(x, t))) => diff_5 + # Differential(x)(Differential(x)(Differential(x)(u(x… => diff_1 + # Differential(x)(Differential(x)(Differential(x)(Dif… => diff_2 + # Differential(x)(u(x, t)) => diff_4 + # Differential(t)(u(x, t)) => diff_3 + + # using HMC algorithm due to convergence, stability, time of training. (refer to mcmc chain plots) + # choice of std for objectives is very important + # pass in Dict_differentials, phystdnew arguments when using the new model + + sol_new = ahmc_bayesian_pinn_pde(pde_system, + discretization; + draw_samples = 150, + bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], phynewstd = [0.4], + phystd = [0.2], l2std = [0.8], param = [Distributions.Normal(2.0, 2)], + priorsNNw = (0.0, 1.0), + saveats = [1 / 100.0, 1 / 100.0], + Dict_differentials = Dict_differentials) + + sol_old = ahmc_bayesian_pinn_pde(pde_system, + discretization; + draw_samples = 150, + bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], + phystd = [0.2], l2std = [0.8], param = [Distributions.Normal(2.0, 2)], + priorsNNw = (0.0, 1.0), + saveats = [1 / 100.0, 1 / 100.0]) + + phi = discretization.phi[1] + xs, ts = [infimum(d.domain):dx:supremum(d.domain) + for (d, dx) in zip(domains, [dx / 10, dt])] + u_real = [[u_analytic(x, t) for x in xs] for t in ts] + + u_predict_new = [[first(pmean(phi([x, t], sol_new.estimated_nn_params[1]))) for x in xs] + for t in ts] + + diff_u_new = [[abs(u_analytic(x, t) - + first(pmean(phi([x, t], sol_new.estimated_nn_params[1])))) + for x in xs] + for t in ts] + + u_predict_old = [[first(pmean(phi([x, t], sol_old.estimated_nn_params[1]))) for x in xs] + for t in ts] + diff_u_old = [[abs(u_analytic(x, t) - + first(pmean(phi([x, t], sol_old.estimated_nn_params[1])))) + for x in xs] + for t in ts] + + unsafe_comparisons(true) + @test all(all, [((diff_u_new[i]) .^ 2 .< 0.8) for i in 1:6]) == true + @test all(all, [((diff_u_old[i]) .^ 2 .< 0.8) for i in 1:6]) == false + + MSE_new = [mean(abs2, diff_u_new[i]) for i in 1:6] + MSE_old = [mean(abs2, diff_u_old[i]) for i in 1:6] + @test (MSE_new .< MSE_old) == [1, 1, 1, 1, 1, 1] + + param_new = sol_new.estimated_de_params[1] + param_old = sol_old.estimated_de_params[1] + α = 1 + @test abs(param_new - α) < 0.2 * α + @test abs(param_new - α) < abs(param_old - α) end diff --git a/lib/BayesianNeuralPDE/test/BPINN_tests.jl b/lib/BayesianNeuralPDE/test/BPINN_tests.jl index 494228275..ef485f4a7 100644 --- a/lib/BayesianNeuralPDE/test/BPINN_tests.jl +++ b/lib/BayesianNeuralPDE/test/BPINN_tests.jl @@ -1,429 +1,429 @@ -@testitem "BPINN ODE I: Without Param Estimation" tags = [:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - using Flux: Flux - - Random.seed!(100) - - linear_analytic = (u0, p, t) -> u0 + sin(2 * π * t) / (2 * π) - linear = (u, p, t) -> cos(2 * π * t) - tspan = (0.0, 2.0) - u0 = 0.0 - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan) - p = prob.p - - # Numerical and Analytical Solutions: testing ahmc_bayesian_pinn_ode() - ta = range(tspan[1], tspan[2], length = 300) - u = [linear_analytic(u0, nothing, ti) for ti in ta] - x̂ = collect(Float64, Array(u) + 0.02 * randn(size(u))) - time = vec(collect(Float64, ta)) - physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] - - # testing points for solve() call must match saveat(1/50.0) arg - ta0 = range(tspan[1], tspan[2], length = 101) - u1 = [linear_analytic(u0, nothing, ti) for ti in ta0] - x̂1 = collect(Float64, Array(u1) + 0.02 * randn(size(u1))) - time1 = vec(collect(Float64, ta0)) - physsol0_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] - - chainlux = Chain(Dense(1, 7, tanh), Dense(7, 1)) - θinit, st = Lux.setup(Random.default_rng(), chainlux) - - fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( - prob, chainlux, draw_samples = 2500) - - alg = BNNODE(chainlux, draw_samples = 2500) - sol1lux = solve(prob, alg) - - # testing points - t = time - # Mean of last 500 sampled parameter's curves[Ensemble predictions] - θ = [vector_to_parameters(fhsamples[i], θinit) for i in 2000:length(fhsamples)] - luxar = [chainlux(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - # --------------------- ahmc_bayesian_pinn_ode() call - @test mean(abs.(x̂ .- meanscurve)) < 0.05 - @test mean(abs.(physsol1 .- meanscurve)) < 0.005 - - #--------------------- solve() call - @test mean(abs.(x̂1 .- pmean(sol1lux.ensemblesol[1]))) < 0.025 - @test mean(abs.(physsol0_1 .- pmean(sol1lux.ensemblesol[1]))) < 0.025 +@testitem "BPINN ODE I: Without Param Estimation" tags=[:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear_analytic = (u0, p, t) -> u0 + sin(2 * π * t) / (2 * π) + linear = (u, p, t) -> cos(2 * π * t) + tspan = (0.0, 2.0) + u0 = 0.0 + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan) + p = prob.p + + # Numerical and Analytical Solutions: testing ahmc_bayesian_pinn_ode() + ta = range(tspan[1], tspan[2], length = 300) + u = [linear_analytic(u0, nothing, ti) for ti in ta] + x̂ = collect(Float64, Array(u) + 0.02 * randn(size(u))) + time = vec(collect(Float64, ta)) + physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] + + # testing points for solve() call must match saveat(1/50.0) arg + ta0 = range(tspan[1], tspan[2], length = 101) + u1 = [linear_analytic(u0, nothing, ti) for ti in ta0] + x̂1 = collect(Float64, Array(u1) + 0.02 * randn(size(u1))) + time1 = vec(collect(Float64, ta0)) + physsol0_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] + + chainlux = Chain(Dense(1, 7, tanh), Dense(7, 1)) + θinit, st = Lux.setup(Random.default_rng(), chainlux) + + fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( + prob, chainlux, draw_samples = 2500) + + alg = BNNODE(chainlux, draw_samples = 2500) + sol1lux = solve(prob, alg) + + # testing points + t = time + # Mean of last 500 sampled parameter's curves[Ensemble predictions] + θ = [vector_to_parameters(fhsamples[i], θinit) for i in 2000:length(fhsamples)] + luxar = [chainlux(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + # --------------------- ahmc_bayesian_pinn_ode() call + @test mean(abs.(x̂ .- meanscurve)) < 0.05 + @test mean(abs.(physsol1 .- meanscurve)) < 0.005 + + #--------------------- solve() call + @test mean(abs.(x̂1 .- pmean(sol1lux.ensemblesol[1]))) < 0.025 + @test mean(abs.(physsol0_1 .- pmean(sol1lux.ensemblesol[1]))) < 0.025 end -@testitem "BPINN ODE II: With Parameter Estimation" tags = [:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - using Flux: Flux - - Random.seed!(100) - - linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) - linear = (u, p, t) -> cos(p * t) - tspan = (0.0, 2.0) - u0 = 0.0 - p = 2 * pi - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, p) - - # Numerical and Analytical Solutions - sol1 = solve(prob, Tsit5(); saveat = 0.01) - u = sol1.u - time = sol1.t - - # BPINN AND TRAINING DATASET CREATION(dataset must be defined only inside problem timespan!) - ta = range(tspan[1], tspan[2], length = 100) - u = [linear_analytic(u0, p, ti) for ti in ta] - x̂ = collect(Float64, Array(u) + 0.2 * randn(size(u))) - time = vec(collect(Float64, ta)) - dataset = [x̂, time] - physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] - - # testing points for solve call(saveat=1/50.0 ∴ at t = collect(eltype(saveat), prob.tspan[1]:saveat:prob.tspan[2] internally estimates) - ta0 = range(tspan[1], tspan[2], length = 101) - u1 = [linear_analytic(u0, p, ti) for ti in ta0] - x̂1 = collect(Float64, Array(u1) + 0.2 * randn(size(u1))) - time1 = vec(collect(Float64, ta0)) - physsol1_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] - - chainlux1 = Chain(Dense(1, 7, tanh), Dense(7, 1)) - θinit, st = Lux.setup(Random.default_rng(), chainlux1) - - fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( - prob, chainlux1, dataset = dataset, draw_samples = 2500, - physdt = 1 / 50.0, priorsNNw = (0.0, 3.0), param = [LogNormal(9, 0.5)]) - - alg = BNNODE(chainlux1, dataset = dataset, draw_samples = 2500, physdt = 1 / 50.0, - priorsNNw = (0.0, 3.0), param = [LogNormal(9, 0.5)]) - - sol2lux = solve(prob, alg) - - # testing points - t = time - # Mean of last 500 sampled parameter's curves(flux and lux chains)[Ensemble predictions] - θ = [vector_to_parameters(fhsamples[i][1:(end-1)], θinit) - for i in 2000:length(fhsamples)] - luxar = [chainlux1(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - # --------------------- ahmc_bayesian_pinn_ode() call - @test mean(abs.(physsol1 .- meanscurve)) < 0.15 - - # ESTIMATED ODE PARAMETERS (NN1 AND NN2) - @test abs(p - mean([fhsamples[i][23] for i in 2000:length(fhsamples)])) < abs(0.35 * p) - - #-------------------------- solve() call - @test mean(abs.(physsol1_1 .- pmean(sol2lux.ensemblesol[1]))) < 8e-2 - - # ESTIMATED ODE PARAMETERS (NN1 AND NN2) - @test abs(p - sol2lux.estimated_de_params[1]) < abs(0.15 * p) +@testitem "BPINN ODE II: With Parameter Estimation" tags=[:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear_analytic = (u0, p, t) -> u0 + sin(p * t) / (p) + linear = (u, p, t) -> cos(p * t) + tspan = (0.0, 2.0) + u0 = 0.0 + p = 2 * pi + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan, p) + + # Numerical and Analytical Solutions + sol1 = solve(prob, Tsit5(); saveat = 0.01) + u = sol1.u + time = sol1.t + + # BPINN AND TRAINING DATASET CREATION(dataset must be defined only inside problem timespan!) + ta = range(tspan[1], tspan[2], length = 100) + u = [linear_analytic(u0, p, ti) for ti in ta] + x̂ = collect(Float64, Array(u) + 0.2 * randn(size(u))) + time = vec(collect(Float64, ta)) + dataset = [x̂, time] + physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] + + # testing points for solve call(saveat=1/50.0 ∴ at t = collect(eltype(saveat), prob.tspan[1]:saveat:prob.tspan[2] internally estimates) + ta0 = range(tspan[1], tspan[2], length = 101) + u1 = [linear_analytic(u0, p, ti) for ti in ta0] + x̂1 = collect(Float64, Array(u1) + 0.2 * randn(size(u1))) + time1 = vec(collect(Float64, ta0)) + physsol1_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] + + chainlux1 = Chain(Dense(1, 7, tanh), Dense(7, 1)) + θinit, st = Lux.setup(Random.default_rng(), chainlux1) + + fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( + prob, chainlux1, dataset = dataset, draw_samples = 2500, + physdt = 1 / 50.0, priorsNNw = (0.0, 3.0), param = [LogNormal(9, 0.5)]) + + alg = BNNODE(chainlux1, dataset = dataset, draw_samples = 2500, physdt = 1 / 50.0, + priorsNNw = (0.0, 3.0), param = [LogNormal(9, 0.5)]) + + sol2lux = solve(prob, alg) + + # testing points + t = time + # Mean of last 500 sampled parameter's curves(flux and lux chains)[Ensemble predictions] + θ = [vector_to_parameters(fhsamples[i][1:(end - 1)], θinit) + for i in 2000:length(fhsamples)] + luxar = [chainlux1(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + # --------------------- ahmc_bayesian_pinn_ode() call + @test mean(abs.(physsol1 .- meanscurve)) < 0.15 + + # ESTIMATED ODE PARAMETERS (NN1 AND NN2) + @test abs(p - mean([fhsamples[i][23] for i in 2000:length(fhsamples)])) < abs(0.35 * p) + + #-------------------------- solve() call + @test mean(abs.(physsol1_1 .- pmean(sol2lux.ensemblesol[1]))) < 8e-2 + + # ESTIMATED ODE PARAMETERS (NN1 AND NN2) + @test abs(p - sol2lux.estimated_de_params[1]) < abs(0.15 * p) end -@testitem "BPINN ODE III" tags = [:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - using Flux: Flux - - Random.seed!(100) - - linear = (u, p, t) -> u / p + exp(t / p) * cos(t) - tspan = (0.0, 10.0) - u0 = 0.0 - p = -5.0 - prob = ODEProblem(linear, u0, tspan, p) - linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) - # SOLUTION AND CREATE DATASET - sol = solve(prob, Tsit5(); saveat = 0.1) - u = sol.u - time = sol.t - x̂ = u .+ (u .* 0.1) .* randn(size(u)) - dataset = [x̂, time] - physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] - - # separate set of points for testing the solve() call (it uses saveat 1/50 hence here length 501) - time1 = vec(collect(Float64, range(tspan[1], tspan[2], length = 501))) - physsol2 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] - - chainlux12 = Chain(Dense(1, 6, tanh), Dense(6, 6, tanh), Dense(6, 1)) - θinit, st = Lux.setup(Random.default_rng(), chainlux12) - - # this a forward solve - fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode( - prob, chainlux12, draw_samples = 500, phystd = [0.01], priorsNNw = (0.0, 10.0)) - - fh_mcmc_chainlux22, fhsampleslux22, fhstatslux22 = ahmc_bayesian_pinn_ode( - prob, chainlux12, dataset = dataset, draw_samples = 500, l2std = [0.02], - phystd = [0.05], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)]) - - alg = BNNODE(chainlux12, dataset = dataset, draw_samples = 500, l2std = [0.02], - phystd = [0.05], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)]) - - sol3lux_pestim = solve(prob, alg) - - # testing timepoints - t = sol.t - #------------------------------ ahmc_bayesian_pinn_ode() call - # Mean of last 500 sampled parameter's curves(lux chains)[Ensemble predictions] - θ = [vector_to_parameters(fhsampleslux12[i], θinit) - for i in 400:length(fhsampleslux12)] - luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve2_1 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - θ = [vector_to_parameters(fhsampleslux22[i][1:(end-1)], θinit) - for i in 400:length(fhsampleslux22)] - luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve2_2 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - @test mean(abs, sol.u .- meanscurve2_1) < 1e-2 - @test mean(abs, physsol1 .- meanscurve2_1) < 1e-2 - @test mean(abs, sol.u .- meanscurve2_2) < 1.5 - @test mean(abs, physsol1 .- meanscurve2_2) < 1.5 - - # estimated parameters(lux chain) - param1 = mean(i[62] for i in fhsampleslux22[400:length(fhsampleslux22)]) - @test abs(param1 - p) < abs(0.5 * p) +@testitem "BPINN ODE III" tags=[:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear = (u, p, t) -> u / p + exp(t / p) * cos(t) + tspan = (0.0, 10.0) + u0 = 0.0 + p = -5.0 + prob = ODEProblem(linear, u0, tspan, p) + linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) + # SOLUTION AND CREATE DATASET + sol = solve(prob, Tsit5(); saveat = 0.1) + u = sol.u + time = sol.t + x̂ = u .+ (u .* 0.1) .* randn(size(u)) + dataset = [x̂, time] + physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] + + # separate set of points for testing the solve() call (it uses saveat 1/50 hence here length 501) + time1 = vec(collect(Float64, range(tspan[1], tspan[2], length = 501))) + physsol2 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] + + chainlux12 = Chain(Dense(1, 6, tanh), Dense(6, 6, tanh), Dense(6, 1)) + θinit, st = Lux.setup(Random.default_rng(), chainlux12) + + # this a forward solve + fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode( + prob, chainlux12, draw_samples = 500, phystd = [0.01], priorsNNw = (0.0, 10.0)) + + fh_mcmc_chainlux22, fhsampleslux22, fhstatslux22 = ahmc_bayesian_pinn_ode( + prob, chainlux12, dataset = dataset, draw_samples = 500, l2std = [0.02], + phystd = [0.05], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)]) + + alg = BNNODE(chainlux12, dataset = dataset, draw_samples = 500, l2std = [0.02], + phystd = [0.05], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)]) + + sol3lux_pestim = solve(prob, alg) + + # testing timepoints + t = sol.t + #------------------------------ ahmc_bayesian_pinn_ode() call + # Mean of last 500 sampled parameter's curves(lux chains)[Ensemble predictions] + θ = [vector_to_parameters(fhsampleslux12[i], θinit) + for i in 400:length(fhsampleslux12)] + luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve2_1 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + θ = [vector_to_parameters(fhsampleslux22[i][1:(end - 1)], θinit) + for i in 400:length(fhsampleslux22)] + luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve2_2 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + @test mean(abs, sol.u .- meanscurve2_1) < 1e-2 + @test mean(abs, physsol1 .- meanscurve2_1) < 1e-2 + @test mean(abs, sol.u .- meanscurve2_2) < 1.5 + @test mean(abs, physsol1 .- meanscurve2_2) < 1.5 + + # estimated parameters(lux chain) + param1 = mean(i[62] for i in fhsampleslux22[400:length(fhsampleslux22)]) + @test abs(param1 - p) < abs(0.5 * p) end -@testitem "BPINN ODE: Translating from Flux" tags = [:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - using Flux: Flux - - Random.seed!(100) - - linear_analytic = (u0, p, t) -> u0 + sin(2 * π * t) / (2 * π) - linear = (u, p, t) -> cos(2 * π * t) - tspan = (0.0, 2.0) - u0 = 0.0 - prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan) - p = prob.p - - # Numerical and Analytical Solutions: testing ahmc_bayesian_pinn_ode() - ta = range(tspan[1], tspan[2], length = 300) - u = [linear_analytic(u0, nothing, ti) for ti in ta] - x̂ = collect(Float64, Array(u) + 0.02 * randn(size(u))) - time = vec(collect(Float64, ta)) - physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] - - # testing points for solve() call must match saveat(1/50.0) arg - ta0 = range(tspan[1], tspan[2], length = 101) - u1 = [linear_analytic(u0, nothing, ti) for ti in ta0] - x̂1 = collect(Float64, Array(u1) + 0.02 * randn(size(u1))) - time1 = vec(collect(Float64, ta0)) - physsol0_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] - chainflux = Flux.Chain(Flux.Dense(1, 7, tanh), Flux.Dense(7, 1)) |> Flux.f64 - fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( - prob, chainflux, draw_samples = 2500) - alg = BNNODE(chainflux, draw_samples = 2500) - @test alg.chain isa AbstractLuxLayer +@testitem "BPINN ODE: Translating from Flux" tags=[:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear_analytic = (u0, p, t) -> u0 + sin(2 * π * t) / (2 * π) + linear = (u, p, t) -> cos(2 * π * t) + tspan = (0.0, 2.0) + u0 = 0.0 + prob = ODEProblem(ODEFunction(linear, analytic = linear_analytic), u0, tspan) + p = prob.p + + # Numerical and Analytical Solutions: testing ahmc_bayesian_pinn_ode() + ta = range(tspan[1], tspan[2], length = 300) + u = [linear_analytic(u0, nothing, ti) for ti in ta] + x̂ = collect(Float64, Array(u) + 0.02 * randn(size(u))) + time = vec(collect(Float64, ta)) + physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] + + # testing points for solve() call must match saveat(1/50.0) arg + ta0 = range(tspan[1], tspan[2], length = 101) + u1 = [linear_analytic(u0, nothing, ti) for ti in ta0] + x̂1 = collect(Float64, Array(u1) + 0.02 * randn(size(u1))) + time1 = vec(collect(Float64, ta0)) + physsol0_1 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] + chainflux = Flux.Chain(Flux.Dense(1, 7, tanh), Flux.Dense(7, 1)) |> Flux.f64 + fh_mcmc_chain, fhsamples, fhstats = ahmc_bayesian_pinn_ode( + prob, chainflux, draw_samples = 2500) + alg = BNNODE(chainflux, draw_samples = 2500) + @test alg.chain isa AbstractLuxLayer end -@testitem "BPINN ODE III: with the new objective" tags = [:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - using Flux: Flux - - Random.seed!(100) - - linear = (u, p, t) -> u / p + exp(t / p) * cos(t) - tspan = (0.0, 10.0) - u0 = 0.0 - p = -5.0 - prob = ODEProblem(linear, u0, tspan, p) - linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) - - # SOLUTION AND CREATE DATASET - sol = solve(prob, Tsit5(); saveat = 0.1) - u = sol.u - time = sol.t - x̂ = u .+ (0.1 .* randn(size(u))) - dataset = [x̂, time] - physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] - - chainlux12 = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) - θinit, st = Lux.setup(Random.default_rng(), chainlux12) - - fh_mcmc_chainlux22, fhsampleslux22, fhstatslux22 = ahmc_bayesian_pinn_ode( - prob, chainlux12, - dataset = dataset, - draw_samples = 500, - l2std = [0.1], - phystd = [0.01], - phynewstd = [0.01], - priorsNNw = (0.0, - 1.0), - param = [ - Normal(-7, 3), - ], estim_collocate = true) - - fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode( - prob, chainlux12, - dataset = dataset, - draw_samples = 500, - l2std = [0.1], - phystd = [0.01], - priorsNNw = (0.0, - 1.0), - param = [ - Normal(-7, 3), - ]) - - # testing timepoints - t = sol.t - #------------------------------ ahmc_bayesian_pinn_ode() call - # Mean of last 100 sampled parameter's curves(lux chains)[Ensemble predictions] - θ = [vector_to_parameters(fhsampleslux12[i][1:(end-1)], θinit) - for i in 400:length(fhsampleslux12)] - luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve2_1 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - θ = [vector_to_parameters(fhsampleslux22[i][1:(end-1)], θinit) - for i in 400:length(fhsampleslux22)] - luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] - luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] - meanscurve2_2 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean - - @test mean(abs.(sol.u .- meanscurve2_2)) < 1e-2 - @test mean(abs.(physsol1 .- meanscurve2_2)) < 1e-2 - @test mean(abs.(sol.u .- meanscurve2_1)) > mean(abs.(sol.u .- meanscurve2_2)) - @test mean(abs.(physsol1 .- meanscurve2_1)) > mean(abs.(physsol1 .- meanscurve2_2)) - - # estimated parameters(lux chain) - param2 = mean(i[62] for i in fhsampleslux22[400:length(fhsampleslux22)]) - @test abs(param2 - p) < abs(0.05 * p) - - param1 = mean(i[62] for i in fhsampleslux12[400:length(fhsampleslux12)]) - @test abs(param1 - p) > abs(0.5 * p) - @test abs(param2 - p) < abs(param1 - p) +@testitem "BPINN ODE III: with the new objective" tags=[:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear = (u, p, t) -> u / p + exp(t / p) * cos(t) + tspan = (0.0, 10.0) + u0 = 0.0 + p = -5.0 + prob = ODEProblem(linear, u0, tspan, p) + linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) + + # SOLUTION AND CREATE DATASET + sol = solve(prob, Tsit5(); saveat = 0.1) + u = sol.u + time = sol.t + x̂ = u .+ (0.1 .* randn(size(u))) + dataset = [x̂, time] + physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)] + + chainlux12 = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) + θinit, st = Lux.setup(Random.default_rng(), chainlux12) + + fh_mcmc_chainlux22, fhsampleslux22, fhstatslux22 = ahmc_bayesian_pinn_ode( + prob, chainlux12, + dataset = dataset, + draw_samples = 500, + l2std = [0.1], + phystd = [0.01], + phynewstd = [0.01], + priorsNNw = (0.0, + 1.0), + param = [ + Normal(-7, 3) + ], estim_collocate = true) + + fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode( + prob, chainlux12, + dataset = dataset, + draw_samples = 500, + l2std = [0.1], + phystd = [0.01], + priorsNNw = (0.0, + 1.0), + param = [ + Normal(-7, 3) + ]) + + # testing timepoints + t = sol.t + #------------------------------ ahmc_bayesian_pinn_ode() call + # Mean of last 100 sampled parameter's curves(lux chains)[Ensemble predictions] + θ = [vector_to_parameters(fhsampleslux12[i][1:(end - 1)], θinit) + for i in 400:length(fhsampleslux12)] + luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve2_1 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + θ = [vector_to_parameters(fhsampleslux22[i][1:(end - 1)], θinit) + for i in 400:length(fhsampleslux22)] + luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)] + luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)] + meanscurve2_2 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean + + @test mean(abs.(sol.u .- meanscurve2_2)) < 1e-2 + @test mean(abs.(physsol1 .- meanscurve2_2)) < 1e-2 + @test mean(abs.(sol.u .- meanscurve2_1)) > mean(abs.(sol.u .- meanscurve2_2)) + @test mean(abs.(physsol1 .- meanscurve2_1)) > mean(abs.(physsol1 .- meanscurve2_2)) + + # estimated parameters(lux chain) + param2 = mean(i[62] for i in fhsampleslux22[400:length(fhsampleslux22)]) + @test abs(param2 - p) < abs(0.05 * p) + + param1 = mean(i[62] for i in fhsampleslux12[400:length(fhsampleslux12)]) + @test abs(param1 - p) > abs(0.5 * p) + @test abs(param2 - p) < abs(param1 - p) end -@testitem "BPINN ODE III: new objective solve call" tags = [:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - using Flux: Flux - - Random.seed!(100) - - linear = (u, p, t) -> u / p + exp(t / p) * cos(t) - tspan = (0.0, 10.0) - u0 = 0.0 - p = -5.0 - prob = ODEProblem(linear, u0, tspan, p) - linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) - - # SOLUTION AND CREATE DATASET - sol = solve(prob, Tsit5(); saveat = 0.1) - u = sol.u - time = sol.t - x̂ = u .+ (0.1 .* randn(size(u))) - dataset = [x̂, time] - - # set of points for testing the solve() call (it uses saveat 1/50 hence here length 501) - time1 = vec(collect(Float64, range(tspan[1], tspan[2], length = 501))) - physsol2 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] - - chainlux12 = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) - θinit, st = Lux.setup(Random.default_rng(), chainlux12) - - alg = BNNODE(chainlux12, - dataset = dataset, - draw_samples = 1000, - l2std = [0.1], - phystd = [0.01], - phynewstd = [0.01], - priorsNNw = (0.0, - 1.0), - param = [ - Normal(-7, 3), - ], numensemble = 200, - estim_collocate = true) - - sol3lux_pestim = solve(prob, alg) - - #-------------------------- solve() call - @test mean(abs.(physsol2 .- pmean(sol3lux_pestim.ensemblesol[1]))) < 1e-2 - - # estimated parameters - param3 = sol3lux_pestim.estimated_de_params[1] - @test abs(param3 - p) < abs(0.05 * p) +@testitem "BPINN ODE III: new objective solve call" tags=[:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + linear = (u, p, t) -> u / p + exp(t / p) * cos(t) + tspan = (0.0, 10.0) + u0 = 0.0 + p = -5.0 + prob = ODEProblem(linear, u0, tspan, p) + linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t)) + + # SOLUTION AND CREATE DATASET + sol = solve(prob, Tsit5(); saveat = 0.1) + u = sol.u + time = sol.t + x̂ = u .+ (0.1 .* randn(size(u))) + dataset = [x̂, time] + + # set of points for testing the solve() call (it uses saveat 1/50 hence here length 501) + time1 = vec(collect(Float64, range(tspan[1], tspan[2], length = 501))) + physsol2 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)] + + chainlux12 = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1)) + θinit, st = Lux.setup(Random.default_rng(), chainlux12) + + alg = BNNODE(chainlux12, + dataset = dataset, + draw_samples = 1000, + l2std = [0.1], + phystd = [0.01], + phynewstd = [0.01], + priorsNNw = (0.0, + 1.0), + param = [ + Normal(-7, 3) + ], numensemble = 200, + estim_collocate = true) + + sol3lux_pestim = solve(prob, alg) + + #-------------------------- solve() call + @test mean(abs.(physsol2 .- pmean(sol3lux_pestim.ensemblesol[1]))) < 1e-2 + + # estimated parameters + param3 = sol3lux_pestim.estimated_de_params[1] + @test abs(param3 - p) < abs(0.05 * p) end -@testitem "BPINN ODE IV: Improvement" tags = [:odebpinn] begin - using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, - AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements - using Flux: Flux - - Random.seed!(100) - - function lotka_volterra(u, p, t) - # Model parameters. - α, δ = p - # Current state. - x, y = u - - # Evaluate differential equations. - dx = (α - y) * x # prey - dy = (x - δ) * y # predator - - return [dx, dy] - end - - # initial-value problem. - u0 = [1.0, 1.0] - p = [1.5, 3.0] - tspan = (0.0, 7.0) - prob = ODEProblem(lotka_volterra, u0, tspan, p) - - # OrdinaryDiffEq.jl solve - dt = 0.1 - solution = solve(prob, Tsit5(); saveat = dt) - - times = solution.t - u = hcat(solution.u...) - x = u[1, :] + (0.5 .* randn(length(u[1, :]))) - y = u[2, :] + (0.5 .* randn(length(u[2, :]))) - dataset = [x, y, times] - - chain = Lux.Chain(Lux.Dense(1, 7, tanh), Lux.Dense(7, 7, tanh), - Lux.Dense(7, 2)) - - alg1 = BNNODE(chain; - dataset = dataset, - draw_samples = 1000, - l2std = [0.5, 0.5], - phystd = [0.5, 0.5], - priorsNNw = (0.0, 1.0), - param = [ - Normal(2, 2), - Normal(2, 2)]) - - alg2 = BNNODE(chain; - dataset = dataset, - draw_samples = 1000, - l2std = [0.5, 0.5], - phystd = [0.5, 0.5], - phynewstd = [1.0, 1.0], - priorsNNw = (0.0, 1.0), - param = [ - Normal(2, 2), - Normal(2, 2)], estim_collocate = true) - - @time sol_pestim1 = solve(prob, alg1; saveat = dt) - @time sol_pestim2 = solve(prob, alg2; saveat = dt) - - unsafe_comparisons(true) - bitvec = abs.(p .- sol_pestim1.estimated_de_params) .> - abs.(p .- sol_pestim2.estimated_de_params) - @test bitvec == ones(size(bitvec)) - - Loss_1 = mean(abs, u[1, :] .- pmean(sol_pestim1.ensemblesol[1])) + - mean(abs, u[2, :] .- pmean(sol_pestim1.ensemblesol[2])) - Loss_2 = mean(abs, u[1, :] .- pmean(sol_pestim2.ensemblesol[1])) + - mean(abs, u[2, :] .- pmean(sol_pestim2.ensemblesol[2])) - - @test Loss_1 > Loss_2 +@testitem "BPINN ODE IV: Improvement" tags=[:odebpinn] begin + using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux, + AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements + using Flux: Flux + + Random.seed!(100) + + function lotka_volterra(u, p, t) + # Model parameters. + α, δ = p + # Current state. + x, y = u + + # Evaluate differential equations. + dx = (α - y) * x # prey + dy = (x - δ) * y # predator + + return [dx, dy] + end + + # initial-value problem. + u0 = [1.0, 1.0] + p = [1.5, 3.0] + tspan = (0.0, 7.0) + prob = ODEProblem(lotka_volterra, u0, tspan, p) + + # OrdinaryDiffEq.jl solve + dt = 0.1 + solution = solve(prob, Tsit5(); saveat = dt) + + times = solution.t + u = hcat(solution.u...) + x = u[1, :] + (0.5 .* randn(length(u[1, :]))) + y = u[2, :] + (0.5 .* randn(length(u[2, :]))) + dataset = [x, y, times] + + chain = Lux.Chain(Lux.Dense(1, 7, tanh), Lux.Dense(7, 7, tanh), + Lux.Dense(7, 2)) + + alg1 = BNNODE(chain; + dataset = dataset, + draw_samples = 1000, + l2std = [0.5, 0.5], + phystd = [0.5, 0.5], + priorsNNw = (0.0, 1.0), + param = [ + Normal(2, 2), + Normal(2, 2)]) + + alg2 = BNNODE(chain; + dataset = dataset, + draw_samples = 1000, + l2std = [0.5, 0.5], + phystd = [0.5, 0.5], + phynewstd = [1.0, 1.0], + priorsNNw = (0.0, 1.0), + param = [ + Normal(2, 2), + Normal(2, 2)], estim_collocate = true) + + @time sol_pestim1 = solve(prob, alg1; saveat = dt) + @time sol_pestim2 = solve(prob, alg2; saveat = dt) + + unsafe_comparisons(true) + bitvec = abs.(p .- sol_pestim1.estimated_de_params) .> + abs.(p .- sol_pestim2.estimated_de_params) + @test bitvec == ones(size(bitvec)) + + Loss_1 = mean(abs, u[1, :] .- pmean(sol_pestim1.ensemblesol[1])) + + mean(abs, u[2, :] .- pmean(sol_pestim1.ensemblesol[2])) + Loss_2 = mean(abs, u[1, :] .- pmean(sol_pestim2.ensemblesol[1])) + + mean(abs, u[2, :] .- pmean(sol_pestim2.ensemblesol[2])) + + @test Loss_1 > Loss_2 end diff --git a/lib/BayesianNeuralPDE/test/runtests.jl b/lib/BayesianNeuralPDE/test/runtests.jl index 75d7889d1..5b775a943 100644 --- a/lib/BayesianNeuralPDE/test/runtests.jl +++ b/lib/BayesianNeuralPDE/test/runtests.jl @@ -13,8 +13,8 @@ const GROUP = lowercase(get(ENV, "GROUP", "all")) using NeuralPDE @info "Running tests with $(RETESTITEMS_NWORKERS) workers and \ - $(RETESTITEMS_NWORKER_THREADS) threads for group $(GROUP)" + $(RETESTITEMS_NWORKER_THREADS) threads for group $(GROUP)" ReTestItems.runtests(NeuralPDE; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]), - nworkers = RETESTITEMS_NWORKERS, - nworker_threads = RETESTITEMS_NWORKER_THREADS, testitem_timeout = 3600) + nworkers = RETESTITEMS_NWORKERS, + nworker_threads = RETESTITEMS_NWORKER_THREADS, testitem_timeout = 3600) diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index bcd78c170..eaa8ea11f 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -28,8 +28,8 @@ using RecursiveArrayTools: DiffEqArray using Reexport: @reexport using RuntimeGeneratedFunctions: RuntimeGeneratedFunctions, @RuntimeGeneratedFunction using SciMLBase: SciMLBase, BatchIntegralFunction, IntegralProblem, NoiseProblem, - OptimizationFunction, OptimizationProblem, ReturnCode, discretize, - isinplace, solve, symbolic_discretize + OptimizationFunction, OptimizationProblem, ReturnCode, discretize, + isinplace, solve, symbolic_discretize using Statistics: Statistics, mean using QuasiMonteCarlo: QuasiMonteCarlo, LatinHypercubeSample using WeightInitializers: glorot_uniform, zeros32 @@ -38,13 +38,13 @@ using Zygote: Zygote # Symbolic Stuff using ModelingToolkit: ModelingToolkit, PDESystem, Differential, toexpr using Symbolics: Symbolics, unwrap, arguments, operation, build_expr, Num, - expand_derivatives + expand_derivatives using SymbolicUtils: SymbolicUtils using SymbolicIndexingInterface: SymbolicIndexingInterface # Needed for the Bayesian Stuff using AdvancedHMC: AdvancedHMC, HMCDA, - NUTS + NUTS using Distributions: Distributions, Distribution, MvNormal, dim, logpdf using LogDensityProblems: LogDensityProblems using MCMCChains: MCMCChains @@ -95,14 +95,14 @@ export DeepGalerkin export neural_adapter export GridTraining, StochasticTraining, QuadratureTraining, QuasiRandomTraining, - WeightedIntervalTraining + WeightedIntervalTraining export build_loss_function, get_loss_function, - generate_training_sets, get_variables, get_argument, get_bounds, - get_numeric_integral, symbolic_discretize, vector_to_parameters + generate_training_sets, get_variables, get_argument, get_bounds, + get_numeric_integral, symbolic_discretize, vector_to_parameters export AbstractAdaptiveLoss, NonAdaptiveLoss, GradientScaleAdaptiveLoss, - MiniMaxAdaptiveLoss + MiniMaxAdaptiveLoss export LogOptions