# python code

"""
Solve initial value problem
for a charge in crossed electric and magnetic fields (homogeneous)
Plot trajectory
"""

from scipy.integrate import odeint
from numpy import array,linspace
import matplotlib.pylab as plt                # for plotting commands


# system of ODEs
def deriv(y,t): # return derivatives of the array y
    """nonlinear ODEs
    charged particle in crossed electric and magnetic fields"""
    return array([ y[2], y[3], y[3]+exb , -y[2]+eyb ])


# - - - - - - - - - - - - - - - - - - - - - - - -
# Main program

# prepare graph
fig=plt.figure(figsize=(6,6))
plt.xlabel(r'$x$',fontsize=20,labelpad=0)
plt.ylabel(r'$y$',fontsize=20,labelpad=0)
xmin = -15.0; xmax = 15.0
ymin = -5.0; ymax = 25.0
dx = 5.0; dy = 5.0
plt.xticks(plt.arange(xmin,xmax+dx,dx),fontsize=14)
plt.yticks(plt.arange(ymin,ymax+dy,dy),fontsize=14)

plt.axis([xmin,xmax,ymin,ymax])


# Parameters of the problem
mass = 1.0
charge = 1.0
bfield = 1.0
ex = -1.0
ey = 0.0

# derived parameters
omega = charge*mass/bfield
if bfield != 0.0:
    exb = ex/bfield
    eyb = ey/bfield
else:
    exb = ex
    eyb = ey
         
# input initial condition
init = input('Give initial condition x,y,vx,vy: ')
yinit = array(init)             # initial values
time_interval = (0.0,20.0)
ntime = 100    # number of points on time interval
time = linspace(time_interval[0],time_interval[1],ntime)

# integrate ODEs
y = odeint(deriv,yinit,time)
print "Final point: ", y[ntime-1,0],",",y[ntime-1,1]

# place coordinates and velocities in separate arrays
# find guiding center coordinates
yp1 = []; vp1 = []    # position and velocity in x
yp2 = []; vp2 = []    # position and velocity in y
rx = []; ry = []      # guiding center coordinates
for iline in range(ntime):
    yp1.append(y[iline,0])
    yp2.append(y[iline,1])
    vp1.append(y[iline,2])
    vp2.append(y[iline,3])
    rx.append(yp1[iline]+vp2[iline])
    ry.append(yp2[iline]-vp1[iline])
  
# plot line
plt.plot(yp1,yp2,'-b')    # plot trajectory
plt.plot(rx,ry,'-r')      # plot guiding center
plt.grid(True)
plt.show()


#savefig('plot1.png', dpi=96)