import math
import numpy
from matplotlib.pyplot import *
import scipy.integrate

alpha=1.0
a1=alpha*alpha/4
a2=alpha/2
vz=1.0
z0=-5.0
def equation(y,t):
    z = z0+vz*t
    B = math.pow(1+z*z,-1.5)
    return [y[1],-a1*y[0]*B*B,-a2*B]
tmax=20.0
te=0.01                         
t = np.arange(start=0,stop=tmax,step=te)
atol=1e-6
rtol=1e-6
y0=[0,0.01,0]
y = scipy.integrate.odeint(equation,y0,t,rtol=rtol,atol=atol)
yt = y.transpose()
r=yt[0]
theta=yt[2]
z=z0+vz*t
figure(figsize=(10,5))
plot(z,r)
xlabel('z')
ylabel('r')
grid()
            

figure(figsize=(10,5))
plot(z,theta)
xlabel('z')
ylabel('theta')
grid()
            

vz=0.8
y0=[0,0.01,0]
y = scipy.integrate.odeint(equation,y0,t,rtol=rtol,atol=atol)
yt = y.transpose()
r=yt[0]
theta=yt[2]
z=z0+vz*t
figure(figsize=(10,5))
plot(z,r)
xlabel('z')
ylabel('r')
grid()
            

vz=2.0
y0=[0,0.01,0]
y = scipy.integrate.odeint(equation,y0,t,rtol=rtol,atol=atol)
yt = y.transpose()
r=yt[0]
theta=yt[2]
z=z0+vz*t
figure(figsize=(10,5))
plot(z,r)
xlabel('z')
ylabel('r')
grid()
            

vz=0.001
tmax=10000.0
te=0.01
t = np.arange(start=0,stop=tmax,step=te)
y0=[0,0.001,0]
y = scipy.integrate.odeint(equation,y0,t,rtol=rtol,atol=atol)
yt = y.transpose()
r=yt[0]
theta=yt[2]
z=z0+vz*t
figure(figsize=(10,5))
plot(z,r)
xlabel('z')
ylabel('r')
grid()
            

figure(figsize=(10,5))
plot(z,theta)
xlabel('z')
ylabel('theta')
grid()
            

def equation(y,t):
    r=y[0]
    r2 = r*r
    z=y[1]
    u = 1+z*z
    B = math.pow(u,-1.5)
    dB = -3*z*math.pow(u,-2.5)
    if ptheta==0:
        dpr = -alpha/4*r*B
        dpz = 0.0
        dtheta = -alpha/2*B
    else:
        dpr = ptheta/(r2*r)-alpha/4*r*B
        dpz = alpha*ptheta*dB
        dtheta = ptheta/r2-alpha/2*B
    return [y[2],y[3],dpr,dpz,dtheta]
                

def trajectoire(z0,r0,pz0,pr0,tmax):
    global ptheta
    B0 = math.pow(1+z0*z0,-1.5)
    ptheta = r0*r0/2*B0
    te = 0.01
    t = np.arange(start=0,stop=tmax,step=te)
    y = scipy.integrate.odeint(equation,[r0,z0,pr0,pz0,0.0],t,rtol=1e-5,atol=1e-5)
    yt = y.transpose()
    r = yt[0]
    z = yt[1]
    plot(z,r)
                

def angle(z0,r0,pz0,pr0,tmax):
    global ptheta
    B0 = math.pow(1+z0*z0,-1.5)
    ptheta = r0*r0/2*B0
    te = 0.01
    t = np.arange(start=0,stop=tmax,step=te)
    y = scipy.integrate.odeint(equation,[r0,z0,pr0,pz0,0.0],t,rtol=1e-5,atol=1e-5)
    yt = y.transpose()
    z = yt[1]
    theta = yt[4]
    plot(z,theta)
                

def projection(z0,r0,pz0,pr0,tmax):
    global ptheta
    B0 = math.pow(1+z0*z0,-1.5)
    ptheta = r0*r0/2*B0
    te = 0.01
    t = np.arange(start=0,stop=tmax,step=te)
    y = scipy.integrate.odeint(equation,[r0,z0,pr0,pz0,0.0],t,rtol=1e-5,atol=1e-5)
    yt = y.transpose()
    r = yt[0]
    theta = yt[4]
    plot(r*numpy.cos(theta),r*numpy.sin(theta))
                

alpha = 1.0
figure(figsize=(8,6))
z0 = -10.0
trajectoire(z0,0.1,1.0,0.01,50.0)
trajectoire(z0,0.0,1.0,0.01,50.0)
axis([z0,10,-0.5,0.5])
xlabel('z')
ylabel('r')
grid()
                

figure(figsize=(8,6))
angle(z0,0.1,1.0,0.01,50.0)
angle(z0,0.0,1.0,0.01,50.0)
xlabel('z')
ylabel('theta')
axis([z0,10,-1.0,0.0])
grid()
                

figure(figsize=(6,12))
subplot(2,1,1)
projection(z0,0.0,1.0,0.01,50.0)
xlabel('x')
ylabel('y')
grid()
axis([-0.5,0.5,-0.5,0.5])
subplot(2,1,2)
projection(z0,0.1,1.0,0.01,50.0)
xlabel('x')
ylabel('y')
axis([-0.5,0.5,-0.5,0.5])
grid()
                

alpha = 1.0
figure(figsize=(8,6))
trajectoire(z0,0.1,1.0,0.005,50.0)
trajectoire(z0,0.1,1.0,0.01,50.0)
trajectoire(z0,0.1,1.0,0.02,50.0)
trajectoire(z0,0.1,1.0,0.04,50.0)
axis([z0,10,-0.5,0.5])
xlabel('z')
ylabel('r')
grid()