#!/usr/bin/env python
# coding: utf-8

# # Homework 2
# **due 09/19/2018 at 11:59pm**
# 
# ## Problem 1
# 1. Write a piece of code to compare the deviation between `ODE_integration_exponential_integrator` and the actual Larmor radius. We call the difference between them the error $e$.
# 2. Vary the number of time steps and the number of integration steps and plot how the error $e$ varies as a function of both.
# 3. What can you conclude?
# 
# ## Problem 2
# When electric fields and magnetic fields are presents, charged particles do not simply gyrate, they move along the electric field direction also. When the electric field is purely perpendicular to B, the charge particle is said to drift, hence the name *ExB drift*. We are going to use `ODE_integration_exponential_integrator` to compute the drift velocity of an electron inside a magnetic field $\vec B=(0,0,1)$ and in an electric field $\vec E=(100,0,0)$.
# 1. Compute analytically the gyration period $t_{Larmor}$
# 2. What is its numerical value?
# 3. Starting with $t_{initial}=0$ and ending at $t_{final}=5\,t_{Larmor}$, find the distance traveled by the electron and infer the velocity from this value.
# 4. Redo 1,2 and 3 for a proton. Is the proton traveling in the same direction?
# 
# ## Problem 3
# If you remember we had a little problem with `OED_integration_E`. If we integrated our ODE for too long, then an electron inside a 1V electric field would eventually go faster than the speed of light. `ODE_integration_exponential_integrator` takes into account relativistic correction and really avoid this problem.
# 1. Using the exact same initial values of position and velocity, find out what is the velocity of the electron after 3s. If it does go faster than the speed of light, increase the number of time steps.
# 
# Let's imagine that you want to build a particle accelerator using a steady electric field of $1kV/m$ to accelerate an electron to $0.999c_{light}$
# 2. How long should the accelerator be?
# 2. How long should it be if you ignore relativistic corrections, i.e. use `OED_integration_E`?

# In[1]:


import math
from math import sqrt as sqrt
import numpy as np
from scipy.linalg import expm
from numpy.linalg import norm 
import matplotlib.pyplot as plt
import matplotlib.colors as colors
#to display inside the notebook!
get_ipython().run_line_magic('matplotlib', 'inline')


# # Solutions
# ## Problem 1
# Let's look first at the impact of the impact of number of time steps and integration steps on `ODE_integration_exponential_integrator`. For that we need to generate a two dimensional grid holding both the number of time steps and the number of iteration steps. We can reuse part of `comp_E` from HW01-Problem 3 to do that. First, let's import all the required functions. 

# In[2]:


def get_E(r,t):
    return np.array([0,0,0])
def get_B(r,t):
    return np.array([0,0,1])
def Lorentz_factor(v):
    v2=norm(v)**2
    c=299792458
    return 1./sqrt(1-v2/c**2)
def EM_field_tensor(E,B):
    c=299792458
    F_mu_nu=np.array([0,E[0]/c,E[1]/c,E[2]/c,E[0]/c,0,-B[2],B[1],E[1]/c,B[2],0,-B[0],E[2]/c,-B[1],B[0],0])
    F_mu_nu.shape=(4,4)
    return F_mu_nu


# Now, we load `ODE_integration_exponential_integrator`. 

# In[3]:


def ODE_integration_exponential_integrator(initial_position,initial_velocity,time_frame,n_steps,n_integrate):
    c=299792458
    loc=np.zeros((3,n_steps)) #3 for x,y,z  
    loc[:,0]=initial_position
    v=initial_velocity
    gamma=Lorentz_factor(v)
    p_old=np.array([m*c,m*v[0],m*v[1],m*v[2]])*gamma
    p_new=np.zeros(4)
    t=0
    dt=(time_frame[1]-time_frame[0])/n_steps
    ds=dt/n_integrate
    for i in range(1,n_steps,1): #solution of the momentum equation
        B=get_B(loc[:,i-1],t)
        E=get_E(loc[:,i-1],t)
        F_mu_nu=EM_field_tensor(E,B)
        p_new=np.matmul(expm(q/(m*gamma)*F_mu_nu*dt),p_old) 
        gamma=p_new[0]/(m*c)
        P=np.zeros((4,4))
        s=0.
        dPds=np.eye(4)*ds/2
        x=np.zeros(4)
        for j in range(0,n_integrate,1): #trapezoidal integration to get the trajectory
            s+=ds
            if (i>0):
                P+=dPds
            dPds=expm(q/(m*gamma)*F_mu_nu*s)*ds/2
            if (i==0):
                P+=dPds
            P+=dPds
        x=np.matmul(P,p_old)/(m*gamma)
        loc[:,i]=x[1:4]+loc[:,i-1] #new position
        p_old=p_new
        t+=dt
    return loc


# We will keep the same initial conditions as the ones used in Ch02. The charged particle will start at $(0,0,0)$ so we can get the Larmor radius easily, by taking the maximum value along the x-direction. We use the gyration period `t_gyration` as our time unit, though it is not required as long as the run the integration by more than several periods.

# In[4]:


m=9e-31
q=-1.6e-19
r_initial=np.array([0,0,0])
v_initial=np.array([1000,0,0])
t_gyration=m/(abs(q)*norm(get_B(r_initial,0)))*(2*math.pi)
r_Larmor=norm(v_initial)*m/(abs(q)*norm(get_B(r_initial,0)))
time_frame=[0,2*t_gyration]


# We now write the loop to compute the error for both integration steps.

# In[5]:


def comp_error(steps,n):
    X=np.zeros(n)
    Y=np.zeros(n)
    err=np.zeros(n)
    k=0
    l=0
    for i in range(steps[0],steps[1],int((steps[1]-steps[0])/n[0])):
        if (k<n[0]): # in case the index goes beyond the array bounds
            for j in range(steps[2],steps[3],int((steps[3]-steps[2])/n[1])):
                if (l<n[1]): #ditto
                    r=ODE_integration_exponential_integrator(r_initial,v_initial,time_frame,i,j)
                    err[k,l]=abs(r_Larmor-r[0,:].max())/r_Larmor*100 #the error is in % of r_Larmor
                    X[k,l]=i
                    Y[k,l]=j
                l+=1
        l=0
        k+=1
    return X,Y,err


# Let's run the code

# In[6]:


steps=[5,30,5,30]
ns=[10,10]
x,y,err=comp_error(steps,ns)


# We plot the error as a function of the number of steps and the number of integration steps. The error is plotted in $\%$ but on the logarithmic scale to be able to see across many orders of magnitude. 

# In[7]:


fig, ax = plt.subplots()

im = plt.imshow(err, cmap=plt.cm.jet, norm=colors.LogNorm(vmin=err.min(), vmax=err.max()),extent=steps)
im.set_interpolation('bilinear')

cb = fig.colorbar(im)
plt.xlabel('Number of steps', rotation=0)
plt.ylabel('Number of integration steps', rotation=90)
cb.ax.set_ylabel('$\log_{10}error(\%)$', rotation=-90)
cb.ax.yaxis.set_label_coords(4.5, 0.5)
plt.show()


# The error is increasing with the number of integration steps (i.e. the number of steps to compute $\vec x$ from $\vec p$), but the error decreases as we increase the number of time steps to compute $\vec p_{new}$ from $\vec p_{old}$.

# ## Problem 2
# We first need to define our new electric field and magnetic field. It was defined for the previous problem with $\vec E=\vec 0$.

# In[21]:


def get_E(r,t):
    return np.array([100,0,0])
def get_B(r,t):
    return np.array([0,0,1])


# Then we define the particle properties and the initial conditions to solve our problem. We will place the particle at the origin $(0,0,0)$ so we can get the distance traveled easily. We will use the Larmor period $$\tau_{Larmor}=\frac{2\pi m}{qB}$$ and print it as the problem asks.

# In[28]:


m=9e-31
q=-1.6e-19
r_initial=np.array([0,0,0])
v_initial=np.array([1000,0,0])
t_Larmor=m/(abs(q)*norm(get_B(r_initial,0)))*(2*math.pi)
print('t_Larmor=',t_Larmor*1e12,'ps')
t_max=5*t_Larmor
time_frame=[0,t_max]


# According to the lecture notes, we expect the particle to drift in the y-direction. using the function
# 
# `r=ODE_integration_exponential_integrator(r_initial,v_initial,time_frame,n_steps,n_integration_steps)`
# 
# So we look only at `r[1,:]`. We have to increase the number of steps until the value of distance settles down.

# In[31]:


for i in range(100,800,100):
    r=ODE_integration_exponential_integrator(r_initial,v_initial,time_frame,i,5)
    n_steps=i
    print(r[1,n_steps-1])


# In[32]:


print('v=',r[1,n_steps-1]/t_max,'m/s')


# According to the lecture notes, the drift velocity should be $v_{E\times B}=\frac{100V/m}{1T}=100m/s$, as we just found numerically.
# 
# We start this problem all over again for a proton

# In[33]:


m=1.6e-27
q=1.6e-19
t_Larmor=m/(abs(q)*norm(get_B(r_initial,0)))*(2*math.pi)
print('t_Larmor=',t_Larmor*1e9,'ns')
t_max=5*t_Larmor
time_frame=[0,t_max]
for i in range(100,600,100):
    r=ODE_integration_exponential_integrator(r_initial,v_initial,time_frame,i,5)
    n_steps=i
    print(r[1,n_steps-1])
print('\n v=',r[1,n_steps-1]/t_max,'m/s')


# We get the right answer here also,i.e. $100m/s$. Indeed, both protons and electrons drift in the same direction. This is an important point, because this drift does not create charge separation. As a result, there is really no force that can counter interact that drift. 
# 
# A direct consequence: if we try to confine a plasma inside a magnetic "bottle" and and electric fields perpendicular to B is present then the particles will drift out of the bottle.

# ## Problem 3
# We now use special relativity to limit the motion of electrons, and clamp their velocity below the speed of light. This is done automatically when we use `ODE_integration_exponential_integrator`. As before, we start with redefining our electromagnetic quantities

# In[13]:


def get_E(r,t):
    return np.array([0,1,0])
def get_B(r,t):
    return np.array([0,0,0])


# Then we set our initial conditions, as defined in Ch02.

# In[14]:


q=-1.6e-19
m=9.1e-31
c=299792458
r_initial=np.array([0,0,0])
v_initial=np.array([0,0,0])
time_frame=[0,3]
n_steps=3000 #this number of steps seems to give good convergence
dt=(time_frame[1]-time_frame[0])/n_steps
r=ODE_integration_exponential_integrator(r_initial,v_initial,time_frame,n_steps,3)


# In[15]:


print('v=',abs((r[1,n_steps-1]-r[1,n_steps-2])/dt/c),'c')


# Indeed we have now an electron which velocity is always below th speed of light. Thank you Prof. Einstein!
# Now let's answer the second question by using a much higher electric field.

# In[16]:


def get_E(r,t):
    return np.array([0,1e3,0])
def get_B(r,t):
    return np.array([0,0,0])


# We could write a code to optimize our answer. Instead we use a faster approach and change the final time to get the right answer.

# In[17]:


time_frame=[0,41.9e-6]
n_steps=3000 #this number of steps seems to give good convergence
dt=(time_frame[1]-time_frame[0])/n_steps
r=ODE_integration_exponential_integrator(r_initial,v_initial,time_frame,n_steps,3)
print('v=',abs((r[1,n_steps-1]-r[1,n_steps-2])/dt/c),'c',' Length =',abs(r[1,n_steps-1]),'m')


# Note that with an electric field field of 1kV, it take $41.9\mu s$ to bring an electron to the $0.999c$. However the accelerating structure has to be $12km$ long! Note that modern accelerators work with electric fields on the order of $1MV/m$. So it would take only 12 m. But modern accelerators try to reach electron speeds on the order of $0.9998c$. So that's still long!

# In[ ]: