[WIP] Bayesian PINN solver for ODE #692

AstitvaAggarwal · 2023-06-06T20:39:43Z

fixes #682

…E.jl into develop

src/turing_MCMC.jl

AstitvaAggarwal · 2023-06-29T13:44:39Z

only changed Project.toml(added some packages), NeuralPDE.jl, turing_MCMC.jl and added 2 new files(BPINN_Tests.jl, advancedHMC_MCMC.jl), the rest is the result of JuliaFormatter.format().

…iors, better training times, accuracy(testing accuracy is bad).

Vaibhavdixit02 · 2023-07-03T17:40:49Z

Why are there formatting changes in this PR, I think it might have been by mistake? Try to see if you can remove those since it makes hard to review relevant files

AstitvaAggarwal · 2023-07-04T04:27:21Z

yeah in commit-978e115 I formatted the whole NeuralPDE directory, ive fixed it now. Sorry for the inconvenience.

…eralisation

>Added support for Lux chains >creating mcmc chains from Samples >more options for arguments

src/advancedHMC_MCMC.jl

…es,acceptance rates.

AstitvaAggarwal · 2023-07-11T18:51:12Z

7mins for finite diff derivative scheme and ~20mins for autodiff, 700 samples for NUTS with multinomial sampling. beyond a certain number of samples the Hamiltonian energy decrease stagnates. Problem specific set of Metrics,integrators,adaptors etc. increases performance, speed(HMC faster by around 5x). I was thinking something like Adaptive weighting of BPINNs might help as well due to the nature of optimization.

,trying to add support for functions(f(u,p,t)) which return a vector and take in a vector for u

src/advancedHMC_MCMC.jl

test/BPINN_Tests.jl

xtalax · 2023-07-26T15:28:24Z

src/advancedHMC_MCMC.jl

+# priors: pdf for W,b + pdf for ODE params
+# lotka specific kwargs here
+
+function ahmc_bayesian_pinn_ode(prob::DiffEqBase.DEProblem, chain,


It would be good to support Float32 for GPU use

.github/workflows/CI.yml

xtalax · 2023-08-07T13:53:48Z

.github/workflows/CI.yml

@@ -13,14 +13,17 @@ jobs:
      fail-fast: false
      matrix:
        group:
+          #fixes 682
+          -ODEBPINN


- ODEBPINN

needs a space. One check runs for each group, not for each testset which is why there are only 5 (should be 6)

xtalax · 2023-08-07T13:55:02Z

src/advancedHMC_MCMC.jl

+    p = Tar.prob.p
+    dt = Tar.physdt
+    if isempty(Tar.dataset[end])
+        t = collect(Float64, Tar.prob.tspan[1]:dt:Tar.prob.tspan[2])


replace all Float64 with eltype(\theta) in case we're running on GPU

Vaibhavdixit02 · 2023-08-14T10:23:50Z

src/advancedHMC_MCMC.jl

+
+# dataset would be (x̂,t)
+# priors: pdf for W,b + pdf for ODE params
+function ahmc_bayesian_pinn_ode(prob::DiffEqBase.DEProblem, chain;


@xtalax does this suit the API in NeuralPDE or should we change this?

Vaibhavdixit02

LGTM, just curious about the API here being suitable which can also be addressed from another PR as per @xtalax preference

ChrisRackauckas · 2023-08-15T23:33:23Z

.github/workflows/CI.yml

+          #fixes 682
+          - ODEBPINN
+


Suggested change

#fixes 682

- ODEBPINN

- ODEBPINN

ChrisRackauckas · 2023-08-15T23:34:47Z

src/advancedHMC_MCMC.jl

@@ -0,0 +1,459 @@
+"""


this docstring needs to move to be right on top of the function, otherwise it won't work

ChrisRackauckas · 2023-08-15T23:35:20Z

src/advancedHMC_MCMC.jl

+    chain::C
+    st::S
+    dataset::Vector{Vector{Float64}}
+    priors::Vector{Distribution}


This is going to be unstable: does it not effect performance?

ive been doing @time and @Btime on the function on the merged struct version and the version you put this review (but uses prob::ODEProblem) the difference dosent change much. In fact on many runs the unstable version performs better.

src/advancedHMC_MCMC.jl

ChrisRackauckas · 2023-08-15T23:39:05Z

src/advancedHMC_MCMC.jl

+                     t[i])
+                   for i in 1:length(out[1, :])]
+    end
+    physsol = hcat(physsol...)


Suggested change

physsol = hcat(physsol...)

physsol = reduce(hcat,physsol)

ChrisRackauckas · 2023-08-15T23:51:24Z

Implementation-wise, I think it generally looks pretty good now. I left a few comments but those are pretty minor.

As for API, I do think that might need to change a little. I think this should be setup as just a solve(prob::ODEProblem, alg::BNNODE(...)) with the details for the fitting process in the BNNODE. Now this would be a bit of a weird ODE solver since the solution cannot be easily represented as an ODESolution, but what we can do is add a .original field to the ODESolution (which exists in most other solution types already anyways so it's a reasonable extension) which then holds the fh_mcmc_chain, fhsamples, fhstats (which we can put in a BPINNStatistics struct or something), and then the "standard" solution computation can be done in advance to construct the sol.u. That would make all of this stuff (https://github.com/SciML/NeuralPDE.jl/pull/692/files#diff-f426df4f589f3ec64dfa582ee30371490608d20f836b399cbfebcb443e832c98R130-R147) just isapprox(sol.u , ...) which makes sense because most users will first want to know "what's the ODE solution"? Something that would be nice to the user would be to represent this as a MonteCarloMeasurements or Measurements type so that plotting and such automatically works with error bounds. As done in NNODE, the t would just be chosen to match the saveat.

Then ahmc_bayesian_pinn_ode can be documented as a lower level API for those who really want to dive in, but solve(prob::ODEProblem, alg::BNNODE(...)) would then work very naturally "just like any other ODE solver", with the extra information always available in sol.original but with the uncertainty information represented in the u. That API would be much nicer than the current post-solution analysis required in the tests.

That all said, creating that higher level API is simply calls to ahmc_bayesian_pinn_ode under the hood and then the post analysis that's exactly done in the tests right now, so I don't think doing this high level API will require any new code. It's just restructuring it to be a bit simpler.

AstitvaAggarwal · 2023-08-16T19:48:45Z

Thank you for the reviews, I'll create a seperate Issue, PR for the High level APIs for both the ODE, PDE solver together. I'll move onto the PDE solver first for now, unless you all would recommend finishing with the APIs first.

ChrisRackauckas · 2023-08-17T03:07:52Z

Finishing the higher level ODE interface would be top priority. This functionality effectively won't exist until that piece is done.

AstitvaAggarwal added 3 commits May 11, 2023 21:46

begin project

d897e3f

Merge branch 'develop' of https://github.com/AstitvaAggarwal/NeuralPD…

b77488a

…E.jl into develop

Started by defining Combined likelihood and priors

907c3ed

Vaibhavdixit02 reviewed Jun 7, 2023

View reviewed changes

src/turing_MCMC.jl Outdated Show resolved Hide resolved

Vaibhavdixit02 reviewed Jun 7, 2023

View reviewed changes

src/turing_MCMC.jl Outdated Show resolved Hide resolved

Vaibhavdixit02 reviewed Jun 7, 2023

View reviewed changes

src/turing_MCMC.jl Outdated Show resolved Hide resolved

AstitvaAggarwal added 4 commits June 12, 2023 00:14

model + custom likelihood + adding tests

c8b5dba

formatted stuff

81c47cc

It works but predictions are superbad

56e7afc

formats, added advancedhmc version of BPINN, BPINN test file

978e115

AstitvaAggarwal added 2 commits June 30, 2023 17:19

Tests -BNN, BPINN,changes in var,full/half logprob

a793b3d

custom logprob in turing_MCMC ,advancedHMC_MCMC-added NN parameter pr…

e817b93

…iors, better training times, accuracy(testing accuracy is bad).

Trying to fix formatting done by mistake in commit-978e115

7980448

AstitvaAggarwal added 3 commits July 6, 2023 12:48

The BPINN ODE solver works well for interpolation, now just needs gen…

b760c60

…eralisation

Performs well in Extrapolation,Interpolation.

7e3a3cd

>Added support for Lux chains >creating mcmc chains from Samples >more options for arguments

autodiff takes ~30mins sample- test file example

31d1f65

Vaibhavdixit02 reviewed Jul 10, 2023

View reviewed changes

src/advancedHMC_MCMC.jl Outdated Show resolved Hide resolved

Less Numerical errors, faster sampling(autodiff),better EBFMI estimat…

8329b8f

…es,acceptance rates.

AstitvaAggarwal closed this Jul 11, 2023

AstitvaAggarwal reopened this Jul 11, 2023

AstitvaAggarwal added 3 commits July 15, 2023 21:09

Inverse parameter estimation, seperate priors for weights and biases

15ec776

,trying to add support for functions(f(u,p,t)) which return a vector and take in a vector for u

Inverse problem gives incorrect parameter esim

c2cce0b

Tests for inverse problem estimation

dfed5ee

Vaibhavdixit02 reviewed Jul 17, 2023

View reviewed changes

src/advancedHMC_MCMC.jl Outdated Show resolved Hide resolved

parameter estimation does not work

f77d04a

Vaibhavdixit02 reviewed Jul 18, 2023

View reviewed changes

test/BPINN_Tests.jl Outdated Show resolved Hide resolved

AstitvaAggarwal added 2 commits August 7, 2023 17:45

Tests run automatically

9a47e36

Test Run automatically attempt-1

c24ed46

xtalax reviewed Aug 7, 2023

View reviewed changes

AstitvaAggarwal added 2 commits August 7, 2023 18:27

minor change

4468534

small change again fr fr

0bf8152

xtalax reviewed Aug 7, 2023

View reviewed changes

AstitvaAggarwal and others added 11 commits August 7, 2023 20:20

changes from reviews-1

b23db49

Removed extra packages used in Tests

08883bd

Merge branch 'SciML:master' into develop

a42b6f7

Test dependancies errors

48dfa66

uses AHMC convienience structs

bca2743

Improved Tests - 1

5357b33

improved Tests - 2

3b35fc3

Fixed the Problem in Test

d467250

improved Tests

f2e502a

pls work

e479c47

This shall work!

4df2d63

Vaibhavdixit02 reviewed Aug 14, 2023

View reviewed changes

Vaibhavdixit02 approved these changes Aug 14, 2023

View reviewed changes

Type checking for dataset, Solver usage made clear

ec24839

ChrisRackauckas reviewed Aug 15, 2023

View reviewed changes

src/advancedHMC_MCMC.jl Outdated Show resolved Hide resolved

ChrisRackauckas reviewed Aug 15, 2023

View reviewed changes

changes from reviews

0428dc5

ChrisRackauckas merged commit a8ef303 into SciML:master Aug 17, 2023
14 of 17 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WIP] Bayesian PINN solver for ODE #692

[WIP] Bayesian PINN solver for ODE #692

AstitvaAggarwal commented Jun 6, 2023

AstitvaAggarwal commented Jun 29, 2023 •

edited

Loading

Vaibhavdixit02 commented Jul 3, 2023

AstitvaAggarwal commented Jul 4, 2023

AstitvaAggarwal commented Jul 11, 2023 •

edited

Loading

xtalax Jul 26, 2023

xtalax Aug 7, 2023 •

edited

Loading

xtalax Aug 7, 2023

Vaibhavdixit02 Aug 14, 2023

Vaibhavdixit02 left a comment

ChrisRackauckas Aug 15, 2023

ChrisRackauckas Aug 15, 2023

ChrisRackauckas Aug 15, 2023

AstitvaAggarwal Aug 17, 2023

ChrisRackauckas Aug 15, 2023

ChrisRackauckas commented Aug 15, 2023

AstitvaAggarwal commented Aug 16, 2023 •

edited

Loading

ChrisRackauckas commented Aug 17, 2023

[WIP] Bayesian PINN solver for ODE #692

[WIP] Bayesian PINN solver for ODE #692

Conversation

AstitvaAggarwal commented Jun 6, 2023

AstitvaAggarwal commented Jun 29, 2023 • edited Loading

Vaibhavdixit02 commented Jul 3, 2023

AstitvaAggarwal commented Jul 4, 2023

AstitvaAggarwal commented Jul 11, 2023 • edited Loading

Choose a reason for hiding this comment

xtalax Aug 7, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Vaibhavdixit02 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ChrisRackauckas commented Aug 15, 2023

AstitvaAggarwal commented Aug 16, 2023 • edited Loading

ChrisRackauckas commented Aug 17, 2023

AstitvaAggarwal commented Jun 29, 2023 •

edited

Loading

AstitvaAggarwal commented Jul 11, 2023 •

edited

Loading

xtalax Aug 7, 2023 •

edited

Loading

AstitvaAggarwal commented Aug 16, 2023 •

edited

Loading