Q:

Solving Vector ODE in Python

Id like to thank whoever answers this in advance, as it is probably a stupid question and a waste of time to answer.
The given code takes the electromagnetic forces on a particle and supposedly plots its trajectory, but I cant figure out how to use rk4 with vectors. Feel like my function is set up wrong.
import numpy as np
import matplotlib.pyplot as plt
from numpy.linalg import inv, det, eig
from numpy.random import randn

def rk4(f, x0, dt):
    tn = 0
    xn = x0

    while True:        
        yield tn,xn

        k1 = dt*f(tn,xn)
        k2 = dt*f(tn+dt/2,xn+k1/2)
        k3 = dt*f(tn+dt/2,xn+k2/2)
        k4 = dt*f(tn+dt,xn+k3)

        xn = xn + (k1+2*k2+2*k3+k4)/6
        tn = tn + dt
#--------------------------------------------------------

def f(t,X):
    x,y,z,xv,yv,zv = X
    v = [xv,yv,zv]
    E = [k*x , k*y , -2*z]
    a = (q/m)*(E+np.cross(v,B))
    Xdot = np.array([v,a])
    return Xdot

q , m , B , k = 1.6e-19, 40, [0,0,1], 10e+4

X0 = np.array([0.001 , 0 , 0.001 , 100 , 0 , 0]) 
solver = rk4(f,X0,10e-7) 

ts = []
Xs = []
for t,X in solver:
    ts.append(t)
    Xs.append(X[0])
    if t > 0.001:
        break

#Xs = np.array(Xs) 
#plt.figure()
#plt.plot(ts,Xs)

Id really appreciate any tips or hints
I suspect the problem stems from the way iv set up the function, as the code crashes once it tries to run it through rk4.

A:

You can not apply vector arithmetic to simple python lists, thus convert the lists first into numpy arrays. You need to return a flat vector, not a matrix, thus use array concatenation to join the two parts.
def f(t,X):
    x,y,z,xv,yv,zv = X
    v = np.array([xv,yv,zv])  #  or v = X[3:]
    E = np.array([k*x , k*y , -2*z]) # or E=k*X[:3]; E[2]=-2*X[2]
    a = (q/m)*(E+np.cross(v,B))
    Xdot = np.concatenate([v,a])
    return Xdot

For the last change see Concatenating two one-dimensional NumPy arrays

