SDE PINN solver

[AstitvaAggarwal]

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Add any other context about the problem here.

[AstitvaAggarwal]

As promised this solver will be completed. I got busy in another project so had to put this on hold. Initial reviews on this would be great. The polynomial chaos expansion for SPDEs will be done later.

[ChrisRackauckas]

That looks to be on the right track.

[AstitvaAggarwal]

would sub_batching as added be a good idea? for sub_batch=2 i get

and for sub_batch=1 i get

the sampled z_i values for the plots are different from the ones used for training, and i use 3 terms of the KKL expansion. the problem being solved is mentioned as in #531 . Further using sol.mean_fit and sol.timepoints got from ensemble mean, timepoints for sub_batch=1 i get

and for sub_batch=2 i get

subbatch=2 seems to perform better (200 samples taken), but i think i may be comparing both cases incorrectly (im following the mean_fit plots).

[AstitvaAggarwal]

So i tried solving again for sub batch=1, 2, 5, the mean_fit error increases as sub batch size increases. The MSE training loss converges smoothly for lower number of sub batches. This could be due a lack of z_i's distribution information being reflected in the dataset (ideal case involves large number of z_i samples). I'm suspecting a NN with probabilistic weights might be better suited to this problem (input t, outputs u), where we have a Random loss function (must choose number of z_i before hand in KKL loss approximation). This would be similar to my BPINN solvers with the exception that we have a stochastic objective. But a doubt arises, would the optimization become too difficult then? or it there any work around to this.

[ChrisRackauckas]

I would think you'd need sample sizes of at least like 100 to be able to smoothly converge?

[AstitvaAggarwal]

Ohh yeah. I'll set up some tests for this. I had tried for like n=10, locally but it was too slow.

[AstitvaAggarwal]

I had tried with sub_batch=2, 5 and 10 but these were not better that the solution of sub_batch=1. Ive added the sub_batch=250 tests, i dont understand why do they fail?. (locally sub_batch=100 solve call takes ~3hrs runtime and fails tests). In case Im testing incorrectly or some code can be sped up, do let me know. Thanks.

[AstitvaAggarwal]

Fixed the issue and made the solver faster ! Tests should pass.

[AstitvaAggarwal]

@ChrisRackauckas All the tests pass. lmk, if something can be improved so we can get this bad boy finally merged.

[AstitvaAggarwal]

How results look like (sol.solution.interp returns a smooth interpolation. Hence, both 2000 test samples case plots appear smooth, In the 1000 test sample case the solver by default returns sol.mean_fit - mean over n=1000 rough path solutions):

[ChrisRackauckas]

This looks correct, though it needs a bit more things tested:

1. Test the strong convergence, i.e. with the same noise process it should give a similar result to the normal SDE solver
2. On the first SDE from https://arxiv.org/abs/1804.04344, what kind of behavior does the solver give?

[AstitvaAggarwal]

@ChrisRackauckas sorry for the delay, the solver was failing miserably for the additive noise tests. Turns out I was performing the initial condition transform twice! Finally figured where the code was going wrong ! I'll be posting the results shortly (yes the new results are pretty good as well)

[AstitvaAggarwal]

@ChrisRackauckas I have added tests for strong and weak convergence in case of both GBM and the Additive Noise Test equation in the testing file. Note that we can also train the SDEPINN via strong or weak convergence, I am defaulting to weak convergence ( basis: strong convergence implies weak convergence, but easier to optimize weak solution over the same n iterations) but the user can always choose if they want to train with a total path-wise loss as well. Let me know what you think on this or want more testing on this part.

Ive used MonteCarloMeasurements.Particles for the returned solutions (BPINNs use this as well). Using this the difference in larger vs smaller sub sample sizes is much clear, the confidence intervals are better for more sub samples (in almost all cases they are equal or within the confidence intervals of the truncated/analytical solutions) . The Diffusion term is learnt much well for more sub samples (makes sense as we basically learn the Wiener basis mapping with solutions better).

Enforcing weak solution while training.

A) Sub samples=1:

B) Sub samples=10:

[AstitvaAggarwal]

And all the related tests pass, requesting reviews.