# python code

"""
Solve initial value problem
for a vortex in external field (harmonic potential)
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([ -k2*y[1], k1*y[0] ])


# - - - - - - - - - - - - - - - - - - - - - - - -
# 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 = -5.0; xmax = 5.0
ymin = -5.0; ymax = 5.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])


# Give parameters of the problem
gamma = 1.0   # vorticity 

k1 = 1.0   # potential parameters
k2 = 2.0


# derived parameters
k1 = k1/gamma
k2 = k2/gamma
         
# input initial condition
init = input('Give initial condition x,y: ')
yinit = array(init)             # initial values
#time_interval = input('Give time interval t0,t1 (same to end): ')
#if time_interval[0] == time_interval[1]: break
time_interval = (0.0,10.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 = []
yp2 = []
for iline in range(ntime):
    yp1.append(y[iline,0])
    yp2.append(y[iline,1])
  
# plot line
#plt.show()
plt.plot(yp1,yp2,'-b')
plt.grid(True)
plt.show()


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