High weak order solvers and adjoint sensitivity analysis for stochastic differential equations

Project summary

In this project, we have implemented new promising tools within the SciML organization which are relevant for tasks such as optimal control or parameter estimation for stochastic differential equations. The high weak order solvers will allow for massive performance advantages for fitting expectations of equations. Instead of automatic differentiation (AD) through the operations of an SDE solver, which scales poorly in memory, one can now use efficient stochastic adjoint sensitivity methods.

Blog posts

The following posts describe the work during the entire period in more detail:

1. GSoC 2020: High weak order SDE solvers and their utility in neural SDEs
2. High weak order SDE solvers

Docs

The documentation of the solvers is available here. Docs with respect to the adjoint sensitivity tools will be available here.

Achievements

Please find below a list of the PRs carried out during GSoC in the different repositories in chronological order within the SciML ecosystem.

StochasticDiffEq.jl

Merged:

• Inplace version of DRI1 scheme
• RI1 method
• tstop fixes for reverse time propagation
• Fixes for EulerHeun scheme
• Speed up the tests for the DRI1 and RI1 schemes
• RI3, RI5, RI6, RDI2WM, RDI3WM, and RDI4WM schemes
• RDI1WM scheme
• Adaptive version of the stochastic Runge-Kutta schemes by embedding
• RS1 and RS2 schemes
• PL1WM scheme
• NON scheme
• Citations for weak methods
• Stochastic improved and modified Euler methods
• COM scheme
• Computationally more efficient NON variant (NON2)

Open:

• Static array tests
• Levy area for non-commutative noise processes

DiffEqSensitivity.jl

Merged:

• Adjoint sensitivities for steady states
• Concrete_solve dispatch for steady state problem
• BacksolveAdjoint for SDEs
• GPU savety for SDE BacksolveAdjoint
• Tests for concrete solve with respect to SDEs
• Alternative differentiation choices (vjps) for noise Jacobian
• Fixes and tests for inplace formulation of BacksolveAdjoint
• Efficient SDE BacksolveAdjoint for scalar noise
• Generalization of the SDE Adjoint for non-diagonal noise processes and diagonal noise processes with mixing terms
• InterpolatingAdjoint for SDEs
• Citations for backsolve, steadystate and interpolation adjoint
• Allow for more general noise processes: replace NoiseGrid by NoiseWrapper
• Checkpointing fix for BacksolveAdjoint in case of ODEs and SDEs
• Cheaper non-diagonal noise tests

Open:

• Support adjoints for SDEs written in the Ito sense

DiffEqNoiseProcess.jl

Merged:

• Multi-dimensional Brownian motion tests
• Bug fix for inplace form of NoiseGrid
• Reversible NoiseWrapper
• Relax the size constraints of the available noise processes
• Generalization of the real-valued white noise process function
• Fix of an extraction issue with NoiseGrid
• Allow NoiseWrapper to start at user-specified time points for interpolating parts of a trajectory
• Extraction and endpoint fixes on NoiseWrapper
• Reversal of SDEs written in the Ito sense

DiffEqGPU.jl

Merged:

• Memory efficient reduction function for ensemble problems

ModelingToolkit.jl

Merged:

• modelingtoolkitize for SDESystem with conversion function between Ito and Stratonovich sense

DiffEqDevTools.jl

Merged:

• NoiseWrapper alternative for analyticless convergence tests of SDE solvers
• test_convergence() dispatch for ensemble simulations
• Work precision set for ensemble problems

DiffEqBase.jl

Merged:

• Fix concrete_solve tests

Future work

There is still a lot that we’d like to do, e.g.,

• Writing up more docs and examples
• Implementing drift-implicit weak stochastic Runge-Kutta solvers
• Finishing the SDE adjoints for the Ito sense
• Implementing a virtual Brownian tree to store the noise processes in O(1) memory
• Setting up an OptimalControl library that allows for easy usage of the new tools within a symbolic interface
• Benchmarking of the new solvers and adjoints

Contributions, suggestions & comments are always welcome! You might like to join our slac channels #diffeq-bridged and #neuralsde to get in touch.

Acknowledgement

I would like to thank my mentors Chris Rackauckas, Moritz Schauer, and Yingbo Ma for their amazing support during this project. It was a great opportunity to work in such an inspiring collaboration and I highly appreciate their detailed feedback. I would also like to thank Christoph Bruder, Niels Lörch, Martin Koppenhöfer, and Michal Kloc for helpful comments on my blog posts. Many thanks to the very supportive julia community and to Google’s open source program for funding this experience!