debugging physic number crunching

physics1.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Aug  8 12:33:35 2017

@author: paul
"""

from math import *
from pylab import plot,show,suptitle,xlabel,ylabel,scatter,grid,legend,title,figure,ylim,xlim,subplots
from numpy import *
from scipy import *
from scipy.optimize import fmin,minimize,fsolve,root,fmin_bfgs,fmin_powell
from pylab import *
import matplotlib.pyplot as plt
from matplotlib import *


def pmqhardedgefast(x,ek):
    def sigma(sigma0,Mtot):
        return np.dot(Mtot,np.dot(sigma0,np.transpose(Mtot)))
    lq1 = 0.006; # length of the quadrupole
    lq2 = 0.003; # length of the quadrupole
    d1 = x[0]
    d2 = x[1]
    d3 = x[2]
    hmax=0.001
    r0 = 10**(-6)
    r0p = 4*10**(-3)
    sigma0=identity(4)
    sigma0[0][0]=r0**2
    sigma0[1][1]= r0p**2
    sigma0[2][2]=r0**2
    sigma0[3][3]=r0p**2
    beamline = array([[0, 0],[d1, 506.9],[d1+lq1, 0], [d1+d2+lq1, -506.9], [d1+d2+2*lq1, 0], [d1+d2+d3+2*lq1 ,537.4], [d1+d2+d3+2*lq1+lq2 ,0], [0.41, 0], [0.41, 0]],float)
    def drift(L):
        return [[1, L ,0 ,0], [0 ,1, 0, 0], [0, 0, 1, L], [0, 0, 0, 1]]
    def Quadlensh(kq,zstep):
        return [[cosh(kq*zstep), 1/kq*sinh(kq*zstep), 0, 0],
                [kq*sinh(kq*zstep), cosh(kq*zstep), 0, 0],
                [0, 0, cos(kq*zstep), 1/kq*sin(kq*zstep)],
                [0, 0, -kq*sin(kq*zstep), cos(kq*zstep)]]
    def Quadlensv(kq,zstep):
        return [[cos(kq*zstep), 1/kq*sin(kq*zstep), 0, 0],
                [-kq*sin(kq*zstep), cos(kq*zstep), 0, 0],
                [0, 0, cosh(kq*zstep), 1/kq*sinh(kq*zstep)],
                [0, 0, kq*sinh(kq*zstep), cosh(kq*zstep)]]
    def Efield(sigma0,hmax):
        a = sqrt(sigma0[0][0])
        b = sqrt(sigma0[2][2])
        kx=sqrt(e0*I/(pi*epsilon0*m0*c0**2*(gamma**3)*a*(a+b)))
        ky=sqrt(e0*I/(pi*epsilon0*m0*c0**2*(gamma**3)*b*(a+b)))
        E=identity(4)
        E[1][0]=kx**2*hmax
        E[3][2]=ky**2*hmax
        return E


    gamma = 1 + ek/0.51099891
    beta = (1-gamma**(-2))**(1/2)
    m0 = 9.10938215*10**(-31)
    e0 = 1.60217646*10**(-19)
    epsilon0 = 8.854*10**(-12)
    c0 = 2.99792458*10**8
    brho = gamma*beta*m0*c0/e0
    Mtot=identity(4);
    i=0
    while i < len(beamline)-1:
        h = beamline[i+1][0]-beamline[i][0]
        l=0
        if beamline[i][1] == 0:
            while l < h:
                Mtot=np.dot(drift(hmax),Mtot)
                l+=hmax
            Mtot=np.dot(drift(h-l),Mtot)
            
        elif beamline[i][1] < 0:
            l=0
            M=identity(4)
            while l < h:
                kq = (-beamline[i][1]/brho)**(1/2)
                M=np.dot(Quadlensh(kq,hmax),M)
                l+=hmax
            Mtot=np.dot(M,Mtot)
            
    
        elif beamline[i][1] > 0:
            l=0
            M=identity(4)
            while l < h:
                kq = (beamline[i][1]/brho)**(1/2)
                M=np.dot(Quadlensv(kq,hmax),M)
                l+=hmax
            #print("check##################")
            #print(type(Quadlensv(kq,h)))
            #print(type(M))
            #print("check##################")
            Mtot=np.dot(M,Mtot)
            
                
        i+=1
    F1=(((Mtot[0][1])))
    F2=(Mtot[2][3])
    F3=(Mtot[0][0]-Mtot[2][2])
    #print(sigma0)
    return F1,F2,F3


x0=[0.00389053825058, 0.00234143819623, 0.00196166116367]
e = arange(3,4.55,.05)
d1a=[]
d2a=[]
d3a=[]

for ek in e:
    ans =root(pmqhardedgefast,x0,args=(ek,),method='lm',options={ 'maxiter': 5000, 'xtol': 10**(-20), 'ftol': 10**(-20)})
    d1a.append(ans.x[0])
    d2a.append(ans.x[1])
    d3a.append(ans.x[2])
    print(ans.x[0],ans.x[1],ans.x[2],ek)
plot(e,d1a,'^c',label = "d1")
plot(e,d2a,'^b',label = "d2")
plot(e,d3a,'^r',label = "d3")
plt.title("PMQ seperations vs Kinetic Energy",fontsize = 20)
plt.ylabel(r'$\mathit{d [cm]}$',fontsize = 30)
plt.xlabel(r'$\mathit{\epsilon_k [Mev]}$',fontsize = 20)
plt.legend(loc = "lower right")
plt.ylim(.0,0.004)
plt.xlim(2.5,5)
grid()
show()