Solving Vector ODE in Python

Question

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.

Look at your full error stack, add a `print` before the line that crashes to debug... find out that `B=1`. How do you expect to take the cross product between a 3-vector and a scalar? — Julien, Dec 04 '18 at 01:00

score 0 · Accepted Answer · answered Dec 04 '18 at 10:16

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

Solving Vector ODE in Python

1 Answers1