scipy odeint with complex initial values

Question

I need to solve a complex-domain-defined ODE system, with complex initial values. scipy.integrate.odeint does not work on complex systems. I rod about cutting my system in real and imaginary part and solve separately, but my ODE system's rhs involves products between dependent variables themselves and their complex conjugates. Haw do I do that? Here is my code, I tried breaking RHS in Re and Im parts, but I don't think the solution is the same as if I wouldn't break it because of the internal products between complex numbers. In my script u1 is a (very)long complex function, say u1(Lm) = f_real(Lm) + 1j* f_imag(Lm).

from numpy import *
from scipy import integrate

def cj(z): return z.conjugate()


def dydt(y, t=0):
    # Notation
    # Dependent Variables
    theta1 = y[0]
    theta3 = y[1]
    Lm = y[2]
    u11 = u1(Lm)
    u13 = u1(3*Lm)
    zeta1 = -2*E*u11*theta1
    zeta3 = -2*E*3*u13*theta3
    # Coefficients
    A0 = theta1*cj(zeta1) + 3*theta3*cj(zeta3)
    A2 = -zeta1*theta1 + 3*cj(zeta1)*theta3 + zeta3*cj(theta1)
    A4 = -theta1*zeta3 - 3*zeta1*theta3
    A6 = -3*theta3*zeta3
    A = - (A2/2 + A4/4 + A6/6)
    # RHS vector components
    dy1dt = Lm**2 * (theta1*(A - cj(A)) - cj(theta1)*A2/2
                     - 3/2*theta3*cj(A2) 
                     - 3/4*cj(theta3)*A4
                     - zeta1)
    dy2dt = Lm**2 * (3*theta3*(A - cj(A)) - theta1*A2/2
                     - cj(theta1)*A4/4
                     - 1/2*cj(theta3)*A6
                     - 3*zeta3)
    dy3dt = Lm**3 * (A0 + cj(A0))
    return array([dy1dt, dy2dt, dy3dt])


t = linspace(0, 10000, 100) # Integration time-step
ry0 = array([0.001, 0, 0.1]) # Re(initial condition)
iy0 = array([0.0, 0.0, 0.0]) # Im(initial condition)
y0 = ry0 + 1j*iy0 # Complex Initial Condition

def rdydt(y, t=0): # Re(RHS)
    return dydt(y, t).real
def idydt(y, t=0): # Im(RHS)
    return dydt(y, t).imag

ry, rinfodict = integrate.odeint(rdydt, y0, t, full_output=True)
iy, iinfodict = integrate.odeint(idydt, y0, t, full_output=True)

The error I get is this TypeError: array cannot be safely cast to required type odepack.error: Result from function call is not a proper array of
floats.

Answer 1

As you've discovered, odeint does not handle complex-valued differential equations, but there is scipy.integrate.complex_ode. complex_ode is a convenience function that takes care of converting the system of n complex equations into a system of 2*n real equations. (Note the discrepancy in the signatures of the functions used to define the equations for odeint and ode. odeint expects f(t, y, *args) while ode (and complex_ode) expect f(y, t, *args).)

A similar convenience function can be created for odeint. In the following code, odeintz is a function that handles the conversion of a complex system into a real system and solving it with odeint. The code includes an example of solving a complex system. It also shows how that system can be converted "by hand" to a real system and solved with odeint. But for a large system, that is a tedious and error prone process; using a complex solver is certainly a saner approach.

import numpy as np
from scipy.integrate import odeint


def odeintz(func, z0, t, **kwargs):
    """An odeint-like function for complex valued differential equations."""

    # Disallow Jacobian-related arguments.
    _unsupported_odeint_args = ['Dfun', 'col_deriv', 'ml', 'mu']
    bad_args = [arg for arg in kwargs if arg in _unsupported_odeint_args]
    if len(bad_args) > 0:
        raise ValueError("The odeint argument %r is not supported by "
                         "odeintz." % (bad_args[0],))

    # Make sure z0 is a numpy array of type np.complex128.
    z0 = np.array(z0, dtype=np.complex128, ndmin=1)

    def realfunc(x, t, *args):
        z = x.view(np.complex128)
        dzdt = func(z, t, *args)
        # func might return a python list, so convert its return
        # value to an array with type np.complex128, and then return
        # a np.float64 view of that array.
        return np.asarray(dzdt, dtype=np.complex128).view(np.float64)

    result = odeint(realfunc, z0.view(np.float64), t, **kwargs)

    if kwargs.get('full_output', False):
        z = result[0].view(np.complex128)
        infodict = result[1]
        return z, infodict
    else:
        z = result.view(np.complex128)
        return z


if __name__ == "__main__":
    # Generate a solution to:
    #     dz1/dt = -z1 * (K - z2)
    #     dz2/dt = L - z2
    # K and L are fixed parameters.  z1(t) and z2(t) are complex-
    # valued functions of t.

    # Define the right-hand-side of the differential equation.
    def zfunc(z, t, K, L):
        z1, z2 = z
        return [-z1 * (K - z2), L - z2] 

    # Set up the inputs and call odeintz to solve the system.
    z0 = np.array([1+2j, 3+4j])
    t = np.linspace(0, 4, 101)
    K = 3
    L = 1
    z, infodict = odeintz(zfunc, z0, t, args=(K,L), full_output=True)

    # For comparison, here is how the complex system can be converted
    # to a real system.  The real and imaginary parts are used to
    # write a system of four coupled equations.  The formulas for
    # the complex right-hand-sides are
    #   -z1 * (K - z2) = -(x1 + i*y1) * (K - (x2 + i*y2))
    #                  = (-x1 - i*y1) * (K - x2 + i(-y2))
    #                  = -x1 * (K - x2) - y1*y2 + i*(-y1*(K - x2) + x1*y2)
    # and
    #   L - z2 = L - (x2 + i*y2)
    #          = (L - x2) + i*(-y2)
    def func(r, t, K, L):
        x1, y1, x2, y2 = r
        dx1dt = -x1 * (K - x2) - y1*y2
        dy1dt = -y1 * (K - x2) + x1*y2
        dx2dt = L - x2
        dy2dt = -y2
        return [dx1dt, dy1dt, dx2dt, dy2dt]

    # Use regular odeint to solve the real system.
    r, infodict = odeint(func, z0.view(np.float64), t, args=(K,L), full_output=True)

    # Compare the two solutions.  They should be the same.  (As usual for
    # floating point calculations, there could be a small difference.)
    delta_max = np.abs(z.view(np.float64) - r).max()
    print "Maximum difference between the complex and real versions is", delta_max


    # Plot the real and imaginary parts of the complex solution.

    import matplotlib.pyplot as plt

    plt.clf()
    plt.plot(t, z[:,0].real, label='z1.real')
    plt.plot(t, z[:,0].imag, label='z1.imag')
    plt.plot(t, z[:,1].real, label='z2.real')
    plt.plot(t, z[:,1].imag, label='z2.imag')
    plt.xlabel('t')
    plt.grid(True)
    plt.legend(loc='best')
    plt.show()

Here's the plot generated by the script:

---

Update

This code has been significantly expanded into a function called odeintw that handles complex variables and matrix equations. The new function can be found on github: https://github.com/WarrenWeckesser/odeintw

Answer 2

I think I found a solution by myself. I'm posting it as anybody would find it useful. It appears that odeint cannot deal with complex numbers. Anyway scipy.integrate.ode does by making use of the 'zvode' integration method.