hgrecco / numbakit-ode

Leveraging numba to speed up ODE integration
Other
68 stars 3 forks source link

complex domain problems #31

Open piperfw opened 2 months ago

piperfw commented 2 months ago

The docstrings of the RungeKutta solvers indicate that they can be applied in the complex domain. However a complex initial state y0 is cast to float by the Solver class (core.py L141). Is integration in the complex domain supported?

maurosilber commented 2 months ago

That's just because we copied the docstrings from SciPy. It doesn't seem that straightforward to directly support complex numbers. If the RHS function is already defined using complex numbers, you could wrapping as in the following example:

import nbkode
import numba
import numpy as np

@numba.njit
def complex_f(t, y):
    return -1j * y

@numba.njit
def f(t, y):
    y = y.view(np.complex_)
    dy = complex_f(t, y)
    return dy.view(np.float_)

t0 = 0
y0 = np.array([1.0 + 1.0j], dtype=np.complex_)
solver = nbkode.RungeKutta45(f, t0, y0.view(np.float_))
t, y = solver.step(n=5)
y = y.view(np.complex_)
piperfw commented 2 months ago

Thanks! I was looking at just splitting up the equations into real and imaginary parts, but for a large complex system your wrapping will be far more convenient.

Happy to PR for the docstrings otherwise feel free to close this issue.