changes from reviews

docs/src/manual/bpinns.md

@@ -1,15 +1,15 @@
 # `BayesianPINN` Discretizer for PDESystems
 
-Using the Bayesian PINNs solvers, we can solve general nonlinear PDEs,ODEs and Also simultaniously perform PDE,ODE parameter Estimation.
+Using the Bayesian PINN solvers, we can solve general nonlinear PDEs, ODEs and also simultaneously perform parameter estimation on them.
 
-Note: The BPINN PDE solver also works for ODEs defined using ModelingToolkit, [ModelingToolkit.jl PDESystem documentation](https://docs.sciml.ai/ModelingToolkit/stable/systems/PDESystem/). Despite this the ODE specific BPINN solver `BNNODE` [refer](https://docs.sciml.ai/NeuralPDE/dev/manual/ode/#NeuralPDE.BNNODE) exists and uses `NeuralPDE.advancedhmc_pinn_ode` at a lower level.
+Note: The BPINN PDE solver also works for ODEs defined using ModelingToolkit, [ModelingToolkit.jl PDESystem documentation](https://docs.sciml.ai/ModelingToolkit/stable/systems/PDESystem/). Despite this, the ODE specific BPINN solver `BNNODE` [refer](https://docs.sciml.ai/NeuralPDE/dev/manual/ode/#NeuralPDE.BNNODE) exists and uses `NeuralPDE.ahmc_bayesian_pinn_ode` at a lower level.
 
 # `BayesianPINN` Discretizer for PDESystems and lower level Bayesian PINN Solver calls for PDEs and ODEs.
 
 ```@docs
 NeuralPDE.BayesianPINN
-NeuralPDE.advancedhmc_pinn_pde
-NeuralPDE.advancedhmc_pinn_ode
+NeuralPDE.ahmc_bayesian_pinn_ode
+NeuralPDE.ahmc_bayesian_pinn_pde
 ```
 
 ## `symbolic_discretize` for `BayesianPINN` and lower level interface.

docs/src/tutorials/low_level_2.md

@@ -85,20 +85,19 @@ function generate_dataset_matrix(domains, dx, dt)
 end
 
 datasetpde = [generate_dataset_matrix(domains, dx, dt)]
-plot(datasetpde[1][:, 2], datasetpde[1][:, 1], title="Dataset from Analytical Solution")
 
-# Add noise to dataset
-datasetpde[1][:, 1] = datasetpde[1][:, 1] .+ randn(size(datasetpde[1][:, 1])) .* 5 / 100 .*
-                      datasetpde[1][:, 1]
-plot!(datasetpde[1][:, 2], datasetpde[1][:, 1])
+# noise to dataset
+noisydataset = deepcopy(datasetpde)
+noisydataset[1][:, 1] = noisydataset[1][:, 1] .+ randn(size(noisydataset[1][:, 1])) .* 5 / 100 .*
+                      noisydataset[1][:, 1]
 
 # Neural network
 chain = Lux.Chain(Lux.Dense(2, 8, Lux.tanh),
     Lux.Dense(8, 8, Lux.tanh),
     Lux.Dense(8, 1))
 
 discretization = NeuralPDE.BayesianPINN([chain],
-    GridTraining([dx, dt]), param_estim = true, dataset = [datasetpde, nothing])
+    GridTraining([dx, dt]), param_estim = true, dataset = [noisydataset, nothing])
 
 @named pde_system = PDESystem(eq,
     bcs,
@@ -120,6 +119,10 @@ sol1 = ahmc_bayesian_pinn_pde(pde_system,
 And some analysis:
 
 ```@example low_level_2
+# Dataset
+plot(datasetpde[1][:, 2], datasetpde[1][:, 1], title="Dataset from Analytical Solution")
+plot!(noisydataset[1][:, 2], noisydataset[1][:, 1])
+
 phi = discretization.phi[1]
 xs, ts = [infimum(d.domain):dx:supremum(d.domain) for (d, dx) in zip(domains, [dx / 10, dt])]
 u_predict = [[first(pmean(phi([x, t], sol1.estimated_nn_params[1]))) for x in xs]

src/PDE_BPINN.jl

@@ -295,7 +295,57 @@ function inference(samples, pinnrep, saveats, numensemble, ℓπ)
     end
 end
 
-# priors: pdf for W,b + pdf for ODE params
+"""
+```julia
+ahmc_bayesian_pinn_pde(pde_system, discretization;
+        draw_samples = 1000,
+        bcstd = [0.01], l2std = [0.05],
+        phystd = [0.05], priorsNNw = (0.0, 2.0),
+        param = [], nchains = 1, Kernel = HMC(0.1, 30),
+        Adaptorkwargs = (Adaptor = StanHMCAdaptor,
+            Metric = DiagEuclideanMetric, targetacceptancerate = 0.8),
+        Integratorkwargs = (Integrator = Leapfrog,), saveats = [1 / 10.0],
+        numensemble = floor(Int, draw_samples / 3), progress = false, verbose = false)               
+```
+## NOTES 
+* Dataset is required for accurate Parameter estimation + solving equations.
+* Returned solution is a BPINNsolution consisting of Ensemble solution, estimated PDE and NN parameters
+  for chosen `saveats` grid spacing and last n = `numensemble` samples in Chain. the complete set of samples
+  in the MCMC chain is returned as `fullsolution`,  refer `BPINNsolution` for more details.
+
+## Positional Arguments
+* `pde_system`: ModelingToolkit defined PDE equation or system of equations.
+* `discretization`: BayesianPINN discretization for the given pde_system, Neural Network and training strategy.
+
+## Keyword Arguments
+* `draw_samples`: number of samples to be drawn in the MCMC algorithms (warmup samples are ~2/3 of draw samples)
+* `bcstd`: Vector of standard deviations of BPINN prediction against Initial/Boundary Condition equations.
+* `l2std`: Vector of standard deviations of BPINN prediction against L2 losses/Dataset for each dependant variable of interest.
+* `phystd`: Vector of standard deviations of BPINN prediction against Chosen Underlying PDE equations.
+* `priorsNNw`: Tuple of (mean, std) for BPINN Network parameters. Weights and Biases of BPINN are Normal Distributions by default.
+* `param`: Vector of chosen PDE's parameter's Distributions in case of Inverse problems.
+* `nchains`: number of chains you want to sample
+
+# AdvancedHMC.jl is still developing convenience structs so might need changes on new releases.
+* `Kernel`: Choice of MCMC Sampling Algorithm object with appropriat args/kwargs (AdvancedHMC.jl implemenations HMC/NUTS/HMCDA).
+* `Adaptorkwargs`: A NamedTuple containing the chosen Adaptor, it's Metric and targetacceptancerate, as follows :
+    * `Adaptor`: https://turinglang.org/AdvancedHMC.jl/stable/
+    * `Metric`: https://turinglang.org/AdvancedHMC.jl/stable/
+    * `targetacceptancerate`: Target percentage(in decimal) of iterations in which the proposals were accepted(0.8 by default)
+* `Integratorkwargs`: A NamedTuple containing the chosen integrator and its keyword Arguments, as follows :
+    * `Integrator`: https://turinglang.org/AdvancedHMC.jl/stable/
+    * `jitter_rate`: https://turinglang.org/AdvancedHMC.jl/stable/
+    * `tempering_rate`: https://turinglang.org/AdvancedHMC.jl/stable/
+* `saveats`: Grid spacing for each independant variable for evaluation of ensemble solution, estimated parameters.
+* `numensemble`: Number of last samples to take for creation of ensemble solution, estimated parameters.
+* `progress`: controls whether to show the progress meter or not.
+* `verbose`: controls the verbosity. (Sample call args in AHMC)
+
+"""
+
+"""
+priors: pdf for W,b + pdf for PDE params
+"""
 function ahmc_bayesian_pinn_pde(pde_system, discretization;
         draw_samples = 1000,
         bcstd = [0.01], l2std = [0.05],

src/advancedHMC_MCMC.jl

@@ -437,11 +437,11 @@ Incase you are only solving the Equations for solution, do not provide dataset
 ## 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`: intial parameter values for BPINN (ideally for multiple chains different initializations preferred)
-* `nchains`: number of chains you want to sample (random initialisation of params by default)
+* `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 predicition against L2 losses/Dataset
-* `phystd`: standard deviation of BPINN predicition against Chosen Underlying ODE System
-* `priorsNNw`: Vector of [mean, std] for BPINN parameter. Weights and Biases of BPINN are Normal Distributions by default
+* `l2std`: standard deviation of BPINN prediction against L2 losses/Dataset
+* `phystd`: standard deviation of BPINN prediction against Chosen Underlying ODE System
+* `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)

src/pinn_types.jl

@@ -166,6 +166,11 @@ BayesianPINN(chain,
                   kwargs...) where {iip}
 ```
 
+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.
+
 ## Positional Arguments
 
 * `chain`: a vector of Flux.jl or Lux.jl chains with a d-dimensional input and a