I’m interested in solving a certain Fokker-Planck equation. The goal is to learn an underlying potential energy function given a time series of the evolution of a distribution.

I’m planning to use the approach of this paper. Briefly, they discretize the distribution in space and set up a master equation describing hops between adjacent grid points where the rate is governed by the difference in energy. This results in physically and thermodynamically consistent evolution of the distribution.

To predict the dynamics at some time t, they use the matrix exponential:

p(x, t) = \exp(R(t-t_0)) p(x, t_0).

This has the nice property of ensuring normalization and positivity given a p(x, t_0) with those properties. This is critically important, as p is a probability density.

That said, I have two concerns about this.

First, the rate matrix R will be very sparse, but it’s exponential will be dense, resulting in a huge blow-up in memory usage and limiting problem sizes.

Second, my understanding is that matrix exponential uses BLAS. Will it be possible to differentiate through this? I.e. is there an adjoint defined?

Alternatively, I could imagine solving the master equation using a standard ODE routine. Is there any way to ensure normalization and positivity? Is there a preferred solver for this type of problem?

Finally, setting up the rate matrix likely involves all sorts of indexing gymnastics, as the discrete grid needs to be flattened into a vector to express the problem in matrix form. Are there any routines to help make this less error-prone?