import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as ode
import mpl_toolkits as mplot3d

q=1.602e-19
m=1.67e-27
r_maa=6371.
nt=20000
dt=0.01
relerr=1.e-9 
abserr=1.e-9

def func(y,t):
 dy=np.zeros_like(y)
 b=[0.,0.,0.]
 re=[y[0]/r_maa,y[1]/r_maa,y[2]/r_maa]
 b=emag(re)
 dy[0]=y[3]
 dy[1]=y[4]
 dy[2]=y[5]
 dy[3]=qm*(y[4]*b[2]-y[5]*b[1])
 dy[4]=qm*(y[5]*b[0]-y[3]*b[2])
 dy[5]=qm*(y[3]*b[1]-y[4]*b[0])
 return dy

def emag(r):
 b=[0.,0.,0.]
 b0=3.07e-05
 x=r[0]
 y=r[1]
 z=r[2]
 rr=np.sqrt(x**2+y**2+z**2)
 b[0]=-3.*b0*z*x/(rr**5)
 b[1]=-3.*b0*z*y/(rr**5)
 b[2]=-b0*(2*z**2-x**2-y**2)/(rr**5)
 return b

y0=[38000.,0.,0.,0.,1.e4,2.e4]
qm=q/m
time=np.arange(0,nt*dt,dt)
res=ode.odeint(func,y0,time,rtol=relerr,atol=abserr)

p=plt.axes(projection="3d")
p.plot(res[:,0],res[:,1],res[:,2])
plt.show()