Added a separate package BayesianNeuralPDE.jl

Project.toml

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

lib/BayesianNeuralPDE/Project.toml

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

src/BPINN_ode.jl → lib/BayesianNeuralPDE/src/BPINN_ode.jl

@@ -1,24 +1,24 @@
 # 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
 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)
 ```
@@ -79,7 +79,7 @@ Kevin Linka, Amelie Schäfer, Xuhui Meng, Zongren Zou, George Em Karniadakis, El
 """
 @concrete struct BNNODE <: NeuralPDEAlgorithm
     chain <: AbstractLuxLayer
-    kernel
+    kernel::Any
     strategy <: Union{Nothing, AbstractTrainingStrategy}
     draw_samples::Int
     priorsNNw::Tuple{Float64, Float64}
@@ -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
 
 """
@@ -149,10 +149,10 @@ contains fields related to that).
 """
 @concrete struct BPINNsolution
     original <: BPINNstats
-    ensemblesol
-    estimated_nn_params
-    estimated_de_params
-    timepoints
+    ensemblesol::Any
+    estimated_nn_params::Any
+    estimated_de_params::Any
+    timepoints::Any
 end
 
 function SciMLBase.__solve(prob::SciMLBase.ODEProblem, alg::BNNODE, args...; dt = nothing,

lib/BayesianNeuralPDE/src/BayesianNeuralPDE.jl

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

src/PDE_BPINN.jl → lib/BayesianNeuralPDE/src/PDE_BPINN.jl

@@ -8,9 +8,9 @@
     names::Tuple
     extraparams::Int
     init_params <: Union{AbstractVector, NamedTuple, ComponentArray}
-    full_loglikelihood
-    L2_loss2
-    Φ
+    full_loglikelihood::Any
+    L2_loss2::Any
+    Φ::Any
 end
 
 function LogDensityProblems.logdensity(ltd::PDELogTargetDensity, θ)
@@ -20,7 +20,8 @@ function LogDensityProblems.logdensity(ltd::PDELogTargetDensity, θ)
                priorlogpdf(ltd, θ) + L2LossData(ltd, θ)
     else
         return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) +
-               priorlogpdf(ltd, θ) + L2LossData(ltd, θ) + ltd.L2_loss2(setparameters(ltd, θ), ltd.phynewstd)
+               priorlogpdf(ltd, θ) + L2LossData(ltd, θ) +
+               ltd.L2_loss2(setparameters(ltd, θ), ltd.phynewstd)
     end
 end
 
@@ -57,11 +58,11 @@ function get_lossy(pinnrep, dataset, Dict_differentials)
     # 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]
+                                   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
@@ -251,13 +252,13 @@ function inference(samples, pinnrep, saveats, numensemble, ℓπ)
 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,8 +305,8 @@ 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],