SDE-PINN inverse + strong/weak + improved base solver & docstrings.

[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

• The Base SDE-PINN solver had the following issues (Fixed them) :

  • was not efficient in input passing to sdephi.
  • Batching performance was abysmal (not actually performing batching).
  • non-batching case was also slow.
  • issues with inner_sde_loss loss construction handling.
  • du_dt (like sdephi and inner_sde_loss) handled only Nested Vector inputs (with nested loops at that).
  • the comments in the code were not clear and very complicated.

• The SDE-PINN solver did not support Inverse solving (additional_loss is inadequate). Needed a new loss specification :

  • Formulated and implemented 2 types of losses :
    • Naive Moment Matching of SDE-PINN outputs with the dataset provided.
    • Moment matching for Distribution got via Euler-Maruyama discretization scheme applied for SDE Parameter Estimation.

• The SDE-PINN solver ONLY supported weak-solving SDEs, now it can solve for the strong solution at n fixed paths :

  • add_rand_coeff_2 for strong solving.
  • add_rand_coeff for weak solving - This shows pretty good results for the expected solution to the SDE, i still need to figure out the mathematics but in my opinion it has to do with the law of large Numbers, Monte Carlo sampling for the Random coefficients basis in the truncated KKL expansion.

• Added fields to the returned solver object :

  • estimated SDE parameters.
  • Exact input path information used during training (strong/weak solves are different)
  • sdephi inputs created from Dataset used for Inverse solves.
  • Random ODE solution object. (contains sdephi, trained parameters etc.) via sol.rode_solution.
  • solution's ensemble_fit and ensemble_inputs used for solution validation.
  • estimated_sol field contains MonteCarloMeasurements.Particles objects created from ensemble_fit. It is a SDE-PINN output sized Vector.

A few plots showing Inverse + strong solving the Geometric Brownian Motion SDE with a dataset consisting of n=15 paths, dt=0.01 in the timespan (0, 1). The solve/training is done on 1 generated strong path (we solve over the truncation, therefore generate a set of random coefficients [z_1, z_2, .....z_i] for each path).
The number of paths can be controlled via sub_batch arg to NNSDE.

(Ran the solve twice)

Estimated Parameters (the solver's loss formulation for the diffusion parameter is sign invariant, therefore might get a mixed up +/- for the final estimated parameter). Here the dataset is as described above (to get values closer to ideal_p increase n), the PINN collocation loss had a discretization dt=0.02. parameter initialization: [0.0, 0.0].

julia> ideal_p
2-element Vector{Float64}:
 1.5
 0.5

julia> sol_2.estimated_params
2-element Vector{Float64}:
 1.2268535217978696
 -0.5335091621614064

julia> sol_1.estimated_params
2-element Vector{Float64}:
 1.226853523016635
 0.5335091597301435

[AstitvaAggarwal]

@ChrisRackauckas requesting reviews.