I’m struggling to find a solution for these coupled differential equations:

dz1/dt = z1(t) H11(t) + z2(t) H12(t)

dz2/dt = z2(t) H22(t) + z1(t) H21(t)

Here, z1(t) and z2(t) are unknown functions that I would like to solve for t, and H11, H12, H21, and H22 are time-dependent coefficients of a (2,2) matrix. These matrices came from a different computation and saved as npy file.

I loaded these coefficients in a NumPy array (let’s call it m) of shape (100, 2, 2). So, to access H11 at time t=15, you’d use m[14, 0, 0] and so on. Anyway, the code works fine when H12, H22 etc. are constants (either real or complex) or H12, H21 etc. have functional forms like e^(-it). But I am having trouble loading them from the array.

Any tips or tricks from the community would be greatly appreciated! I found complex_ode method in scipy but it don’t think it is for coupled equations. I tried to use odeintw in the following code.

My code:

from odeintw import odeintw

dt = 0.1
freq = 100
num_time_steps = 100

# Define the coupled differential equations
def Haa(t): 
    return m[int(t)//freq, 0, 0]

def Hbb(t): 
    return m[int(t)//freq, 1, 1]

def Hab(t): 
    return m[int(t)//freq, 0, 1]

def Hba(t): 
    return m[int(t)//freq, 1, 0]

def equations(z, t, Haa, Hbb, Hab, Hba):
    z1, z2 = z
    dz1dt = (z1 * Haa(t) + z2 * Hab(t))
    dz2dt = -(z1 * Hba(t) + z2 * Hbb(t))
    return [dz1dt, dz2dt]

# Time points

t = [step * freq * dt for step in range(num_time_steps)]

# Initial conditions
z0 = [0, 1]  # Initial values of z1 and z2

# Solve ODE
solution = odeintw(equations, z0, t, args=(Haa, Hbb, Hab, Hba))

# Extract solutions
z1, z2 = solution[:, 0], solution[:, 1]

I tried odeint and presented my code.