SciML / OrdinaryDiffEq.jl

High performance ordinary differential equation (ODE) and differential-algebraic equation (DAE) solvers, including neural ordinary differential equations (neural ODEs) and scientific machine learning (SciML)
https://diffeq.sciml.ai/latest/
Other
562 stars 211 forks source link

Pseudospectral method for semilinear PDE (Split-Step Fourier, RKIP) #2501

Open Azercoco opened 1 month ago

Azercoco commented 1 month ago

Hi,

I am doing a PhD in nonlinear optics in which we encounter several semilinear parabolic PDE like the Nonlinear Schrödinger Equation (NLSE) and its variant.

Most of these PDEs are semilinear meaning they can be expressed as

$$\partial_t y = \hat{D}(\omega)y + N(y)$$ where $\hat{D}$ is a constant stiff linear operator but which can be easily computed in Fourier space. This allows to create efficient integrator as $\exp(A \times dt)$ can be easily computed.

For now, in the available algorithms, only exponential integrator (ETDRK) are able to take profit of this structure but several other algorithms exist:

The idea is to make the approximation $\exp( h (\hat{D} + N) ) \approx \exp( h \hat{D}) \exp(h N) + \mathcal{O}(h^2)$. For the NLSE, $exp( h \hat{D})$ is easily computed in the spectral/frequency domain while $\exp(h N)$ can be easily computed in the time domain.

There is numerous improvement of this algorithm in the literature with adaptative time stepping, higher order ...

The idea is to make the transformation $y_I(t) = exp(t\hat{D})y$ then after transformation the equation become:

$$\partial_t y_I = N_I(t) y_t$$ with $N_I(t) = exp(t\hat{D}) N exp(-t\hat{D})$. If we write a single Runge-Kutta step, we can see that we actually only compute $exp(h\hat{D})$. This allows to use any explicit Runge-Kutta method with their order and adaptive time stepping.

These algorithms could be really useful for fast integration of semilinear PDE with adaptative time stepping. Because they use FTTs, they also greatly benefit from GPU acceleration.

**Implementation difficulties

Split-Step Fourier method could be implemented within a more generalised Strang splitting for Split ODE Problem with possibility of using both $exp(f_1)$ and $exp(f_2)$

For RKIP, we could define the transformed problem and use any RK algorithm at the loss of efficiency as we would need to re-compute $ exp(t\hat{D})$ at each step. By writing a custom implementation, we could only compute and cache $exp(h\hat{D})$ once.

For adaptative stepping, the issue is that computing $exp(h\hat{D})$ for different values of $h$ can be costly. The solution is to only allow a discrete set of steps values $h_i$ (for e.g $h_i = 0.001 \times 10^{i/10}$ and compute the $exp(h_i\hat{D})$ when needed and cache them.

I have implemented both methods both in Python and Julia with GPU support and would be more than happy to open a PR for the code but I am unsure how the best way of implementing them in the package (Should I create a custom problem definition or reuse SplitODE ?)

Other Implementations

The package FourierFlow.jl was actually made for representing this kind of problem but actually lacks good adaptative time stepping and the only exponential integrator is ETDRK4 (which is not adaptative).

References:

ChrisRackauckas commented 1 month ago

You can already do this with OrdinaryDiffEq.jl integrators as shown in the benchmarks like https://docs.sciml.ai/SciMLBenchmarksOutput/v0.2/MOLPDE/allen_cahn_spectral_wpd/. So I'm not understanding the question?

Azercoco commented 1 month ago

Thanks for the example. I did not saw it.

The interest of RKIP methods and Split-Step are mostly their performance when implemented with FFTs. RKIP can moreover give adaptive time stepping, which ETDRK can not (actually they have been recent paper on adaptive time stepper for ETDRK but not of these have been implemented in any libraries yet).

ChrisRackauckas commented 1 month ago

There's already many split step schemes, many of which are adaptive: https://docs.sciml.ai/DiffEqDocs/latest/solvers/split_ode_solve/#OrdinaryDiffEq.jl.

Azercoco commented 1 month ago

The adaptative methods for split problem does not seems to take profit of semilinear structure of the problem. Currently, for a semilinear problem, we have the choice between:

The methods I am proposing are:

From my personal benchmark, for a fixed timestep, ETDRK methods slightly outperform RKIP which themselves outperform Split-Step (but this can change with the problem specificities).

The strength of the RKIP come from the adaptative time step. With it, it greatly outperform both other methods.

ChrisRackauckas commented 4 weeks ago

Makes sense, yes we should add it and try it out.