import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
def rk4(x,t,tau,derivsRK,qm,B):
half_tau = 0.5*tau
F1 = derivsRK(x,t,qm,B)
t_half = t + half_tau
xtemp = x + half_tau*F1
F2 = derivsRK(xtemp,t_half,qm,B)
xtemp = x + half_tau*F2
F3 = derivsRK(xtemp,t_half,qm,B)
t_full = t + tau
xtemp = x + tau*F3
F4 = derivsRK(xtemp,t_full,qm,B)
xout = x + tau/6.*(F1 + F4 + 2.*(F2+F3))
return xout
def magrk(s,t,qm,B):
air=0.0
r = np.array([s[0], s[1], s[2]])
v = np.array([s[3], s[4], s[5]])
accel = qm*np.cross(v,B)-air*np.linalg.norm(v)*v
deriv = np.array([v[0], v[1], v[2], accel[0], accel[1], accel[2]])
return deriv
def Bfield(B0,Br,r):
L=10.
Bz=B0*(2.+np.cos(2.*np.pi/L*r[2]))
B1=np.array([0.,0.,Bz])
rho=np.sqrt(r[0]**2+r[1]**2)
b2=np.array([r[0],r[1],0.])
B2=Br*rho*np.sin(2.*np.pi/L*r[2])*b2
B=B1+B2
return B
def magline(L,B0,Br):
fig = plt.figure()
axB = Axes3D(fig)
NF=100; dr=L/NF; Nth=20; dth=2.*np.pi/Nth
for k in range(2):
rh=0.5*(k+1)
for j in range(Nth+1):
z=0.
th=dth*j
x=rh*np.cos(th)
y=rh*np.sin(th)
r=[x,y,z]
Px=[]; Py=[]; Pz=[]; Pr=[]
for i in range(NF+1):
Px.append(r[0]); Py.append(r[1]); Pz.append(r[2]); Pr.append(r)
B=Bfield(B0,Br,r)
r+=dr*B/np.linalg.norm(B)
if(k==0): axB.plot_wireframe(Px,Py,Pz)
if(k==1): axB.plot_wireframe(Px,Py,Pz,color='red')
axB.set_title('magnetic field line')
return
S7='%12.4f'*7; S7s='%12s'*7; Sline='-'*80; S6='%12.4f'*6
S10='%10.5f'*10
q=1.; m=1.; qm=q/m; x0=0.; z0=3.0; v0x=1.; B0=2.0; Br=B0/3.
#q=1.6e-19; m=9.1e-31; qm=q/m; x0=0.; v0x=1.e6; B0=1.e-5
pp=2.*np.pi*m/q/B0
L=10.; Nz=10; dz=L/Nz; Nth=24; dth=2.*np.pi/Nth
magline(L,B0,Br)
time = 0.0; NC=0; NR=0
nstep=2000000
tau=0.001*pp
#print (S7s %('time','x','y','z','vx','vy','vz'))
#print Sline
r0 = np.array([x0,0.,z0])
v0 = np.array([v0x,0.,0.0*v0x])
K0=0.5*m*np.linalg.norm(v0)**2
r=r0; v=v0
B=Bfield(B0,Br,r)
state = np.array([r[0],r[1],r[2],v[0],v[1],v[2]])
time=0.; timez=0; Ki=0.5*m*np.linalg.norm(v)**2
Pt=[]; Px=[]; Py=[]; Pz=[]; Pvz=[]
for istep in range(nstep):
K=0.5*m*np.linalg.norm(v)**2
if(time > 125.*pp): break
K=0.5*m*np.linalg.norm(v)**2
Pt.append(time); Px.append(r[0]); Py.append(r[1]); Pz.append(r[2])
Pvz.append(v[2])
y1=r[1]; vy1=v[1]; vz1=v[2]
B=Bfield(B0,Br,r)
state = rk4(state,time,tau,magrk,qm,B)
r = np.array([state[0], state[1], state[2]])
v = np.array([state[3], state[4], state[5]])
if(vy1*v[1] < 0. and v[0] > 0.):
NC+=1
LastStep=istep
period=time
print 'NC=',NC,istep,time/pp
if(vz1*v[2] < 0. and v[2] > 0.):
NR+=1
Ker=np.abs(K-K0)/K0
print NR,'Complete a trip:K=',tau/pp,K,Ker
print S10 %(tau/pp,time/pp,r[0],r[1],r[2],v[0],v[1],v[2],K,Ker)
if(NR==3): break
time+=tau
fig = plt.figure()
ax = Axes3D(fig)
ax.plot_wireframe(Px,Py,Pz)
ax.set_title('trajectory of the changed particle in magnetic bottle')
plt.figure()
plt.plot(Px,Py)
plt.axis('equal')
plt.title('trajectory on xy-plane')
plt.figure()
plt.plot(Pt,Pz,label='z(t)')
plt.plot(Pt,Px,label='x(t)')
plt.plot(Pt,Py,label='y(t)')
sBr=str(Br)
plt.title('$\\tau=$'+str(tau/pp)+' z0='+str(z0)+' L='+str(L)+'\nq='+str(q)+ \
' m='+str(m)+' B0='+str(B0)+' Br='+sBr[:5])
plt.legend()
plt.figure()
plt.plot(Pt,Pvz)
plt.xlabel('t')
plt.ylabel('$v_z$')
plt.show()