Solve ode in python with complex matrix as initial value

Question

I have a von Neumann equation which looks like: dr/dt = - i [H, r], where r and H are square matricies of complex numbers and I need to find r(t) using python script.

Is there any standart instruments to integrate such equations?

When I was solving another aquation with a vector as initial value, like Schrodinger equation: dy/dt = - i H y, I used scipy.integrate.ode function ('zvode'), but trying to use the same function for von Neumann equation gives me the following error:

**scipy/integrate/_ode.py:869: UserWarning: zvode: Illegal input detected. (See printed message.)
ZVODE--  ZWORK length needed, LENZW (=I1), exceeds LZW (=I2)
self.messages.get(istate, 'Unexpected istate=%s' % istate))
  In above message,  I1 =        72   I2 =        24**

Here is the code:

def integrate(r, t0, t1, dt):
  e = linspace(t0, t1, (t1 - t0) / dt + 10)
  g = linspace(t0, t1, (t1 - t0) / dt + 10)
  u = linspace(t0, t1, (t1 - t0) / dt + 10)
  while r.successful() and r.t < t1:
    r.integrate(r.t + dt)
    e[r.t / dt] = abs(r.y[0][0]) ** 2
    g[r.t / dt] = abs(r.y[1][1]) ** 2
    u[r.t / dt] = abs(r.y[2][2]) ** 2
  return e, g, u


# von Neumann equation's
def right_part(t, rho):
  hamiltonian = (h / 2) * array(
    [[delta, omega_s, omega_p / 2.0 * sin(t * w_p)],
    [omega_s, 0.0, 0.0],
    [omega_p / 2.0 * sin(t * w_p), 0.0, 0.0]],
    dtype=complex128)
  return (dot(hamiltonian, rho) - dot(rho, hamiltonian)) / (1j * h)


def create_integrator():
  r = ode(right_part).set_integrator('zvode', method='bdf', with_jacobian=False)
  psi_init = array([[1.0, 0.0, 0.0],
                   [0.0, 0.0, 0.0],
                   [0.0, 0.0, 0.0]], dtype=complex128)
  t0 = 0
  r.set_initial_value(psi_init, t0)
  return r, t0


def main():
  r, t0 = create_integrator()
  t1 = 10 ** -6
  dt = 10 ** -11
  e, g, u = integrate(r, t0, t1, dt)

main()

Answer 1

I have created a wrapper of scipy.integrate.odeint called odeintw that can handle complex matrix equations such as this. See How to plot the Eigenvalues when solving matrix coupled differential equations in PYTHON? for another question involving a matrix differential equation.

Here's a simplified version of your code that shows how you could use it. (For simplicity, I got rid of most of the constants from your example).

import numpy as np
from odeintw import odeintw


def right_part(rho, t, w_p):
    hamiltonian = (1. / 2) * np.array(
        [[0.1, 0.01, 1.0 / 2.0 * np.sin(t * w_p)],
        [0.01, 0.0, 0.0],
        [1.0 / 2.0 * np.sin(t * w_p), 0.0, 0.0]],
        dtype=np.complex128)
    return (np.dot(hamiltonian, rho) - np.dot(rho, hamiltonian)) / (1j)


psi_init = np.array([[1.0, 0.0, 0.0],
                     [0.0, 0.0, 0.0],
                     [0.0, 0.0, 0.0]], dtype=np.complex128)


t = np.linspace(0, 10, 101)
sol = odeintw(right_part, psi_init, t, args=(0.25,))

sol will be a complex numpy array with shape (101, 3, 3), holding the solution rho(t). The first index is the time index, and the other two indices are the 3x3 matrix.