PyEPIC/beambeam.py at master · yuehao/PyEPIC · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
__author__ = 'yuehao'
import scipy.special as spefunc
import scipy.interpolate as itp
import numpy as np
import cmath as math

def beambeam_Gaussian_grid(sigmax, sigmay, rangex=10.0, rangey=10.0, epsilon=1e-3, gridpoints=10, sdds_file=None):

    if abs(sigmax-sigmay)/(sigmax+sigmay) < epsilon:
        xlist=np.linspace(-rangex*sigmax, rangex*sigmax, num=int(rangex*gridpoints)+1, endpoint=True)
        ylist=np.linspace(-rangey*sigmay, rangey*sigmay, num=int(rangey*gridpoints)+1, endpoint=True)
        xlist,ylist=np.meshgrid(xlist,ylist)
        sigr=(sigmax+sigmay)/2.0
        rabs=np.sqrt(xlist*xlist+ylist*ylist)
        mask=(rabs==0)
        invmask=np.logical_not(mask)
        eabs=np.zeros_like(rabs)
        ex=np.zeros_like(rabs)
        ey=np.zeros_like(rabs)
        eabs[invmask]=(1-np.exp(-rabs[invmask]*rabs[invmask]/2/sigr/sigr))/rabs[invmask]/2/np.pi
        ex[invmask]=eabs[invmask]*xlist[invmask]/rabs[invmask]
        ey[invmask]=eabs[invmask]*ylist[invmask]/rabs[invmask]
        eabs[mask]=0
        ex[mask]=0
        ey[mask]=0

    else:
        xlist=np.linspace(-rangex*sigmax, rangex*sigmax, num=int(rangex*gridpoints)+1, endpoint=True)
        ylist=np.linspace(-rangey*sigmay, rangey*sigmay, num=int(rangey*gridpoints)+1, endpoint=True)

        spefunc.errprint(1)
        if sigmax<sigmay:
            su=sigmay
            sv=sigmax
            ulist=ylist
            vlist=xlist
            ulist,vlist=np.meshgrid(ulist,vlist)
        else:
            su=sigmax
            sv=sigmay
            ulist=xlist
            vlist=ylist
            ulist,vlist=np.meshgrid(ulist,vlist)

        xmask=ulist>0
        ymask=vlist>0
        xmask=xmask*2.0-1.0
        ymask=ymask*2.0-1.0
        sqrtsigma=math.sqrt(2*su*su-2*sv*sv)
        w1=spefunc.wofz((np.abs(ulist)+1j*np.abs(vlist))/sqrtsigma)
        w2=spefunc.wofz((np.abs(ulist)*sv/su+1j*np.abs(vlist)*su/sv)/sqrtsigma)
        expterm=np.exp(-ulist*ulist/2.0/su/su-vlist*vlist/2.0/sv/sv)
        ecompx=(w1-expterm*w2)/math.sqrt(np.pi)/2.0/1.0j/sqrtsigma
        eu=np.real(ecompx)
        ev=-np.imag(ecompx)
        if sigmax<sigmay:
            ex=ev*xmask
            ey=eu*ymask
            xlist=vlist
            ylist=ulist
        else:
            ex=eu*xmask
            ey=ev*ymask
            xlist=ulist
            ylist=vlist
    #if sdds_file is not None:
    #    import SDDSIO.sdds as sdds
    #    output=sdds.SDDS(0)
    #    output.columnName=['x','y','Fx','Fy']
    #    output.columnDefinition[['', 'm', 'x', '', output.SDDS_DOUBLE, 0L],
    #                            ['', 'm', 'y', '', output.SDDS_DOUBLE, 0L],
    #                            ['', '', 'Fx', '', output.SDDS_DOUBLE, 0L],
    #                            ['', '', 'Fy', '', output.SDDS_DOUBLE, 0L],
    #    ]
    #    output.columnData[[xlist.tolist(),],[ylist.tolist(),][ex.tolist(),][ey.tolist(),]]
    #    output.save(sdds_file)
    return xlist,ylist,ex,ey

def Gaussian_beam_beam(sigmax, sigmay, xlist, ylist):
    sigxy=np.sqrt(2*(sigmax*sigmax-sigmay*sigmay))


    z1=(xlist+1j*ylist)/sigxy
    z2=(xlist*sigmay/sigmax+1j*ylist*sigmax/sigmay)/sigxy
    expterm=np.exp(-xlist*xlist/sigmax/sigmax/2-ylist*ylist/sigmay/sigmay/2)
    result=spefunc.wofz(z1)-expterm*spefunc.wofz(z2)
    nEy=result.real
    nEx=result.imag

    return nEx,nEy


def flat_Gaussian_beam_beam(sigmax, sigmay, xlist, ylist, order=1):
    r=sigmay/sigmax*order;

    sqrtpi=np.sqrt(np.pi)
    xtilde=xlist/np.sqrt(2.0)/sigmax
    ytilde = ylist / np.sqrt(2.0) / sigmay
    xt2= xtilde*xtilde
    yt2 = ytilde * ytilde
    xy=xtilde*ytilde

    expx=np.exp(-xt2)
    expy=np.exp(-yt2)

    dawsonx=spefunc.dawsn(xtilde)
    erfy=spefunc.erf(ytilde)

    nEy=expx*erfy-2*r*ytilde/sqrtpi+4*r*xy*dawsonx/sqrtpi
    nEx=2*dawsonx/sqrtpi-2*r*xy*expx*erfy-2*r*xtilde*expx*expy/sqrtpi
    nEx+=r*r*(2*yt2*(dawsonx+xtilde*(1-2*xtilde*dawsonx))+xtilde-xt2*dawsonx)/sqrtpi

    return nEx,nEy


if __name__=='__main__':
    #import matplotlib.pyplot as plt
    #from mpl_toolkits.mplot3d.axes3d import Axes3D
    #from matplotlib import cm
    #xl,yl,ex,ey=beambeam_Gaussian_grid(1.e-4, 3.e-4)

    #exit()
    #fig=plt.figure(figsize=plt.figaspect(0.25))
    #ax=fig.add_subplot(1,2,1,projection='3d')
    # ax.set_zlim(0,300000)
    #surf=ax.plot_surface(xl,yl,ex*ex+ey*ey,rstride=1,cstride=1,cmap=cm.coolwarm, linewidth=0, antialiased=False)
    #fig.colorbar(surf, shrink=0.5, aspect=10)
    #fig,ax=plt.subplots()
    #ax.set_xlim([0, 1e-3])
    #ax.set_ylim([0, 1000])
    #ax.plot(yl,ex)
    #plt.show()
    pass