atcollab / at

Accelerator Toolbox
Apache License 2.0
48 stars 31 forks source link

modification to get_lifetime #622

Closed bdalena closed 1 year ago

bdalena commented 1 year ago

Hello,

In order to compute the local Piwinski scattering rate we would like to have the function get_lifetime modified as follow:

def get_lifetime(ring, emity, bunch_curr, emitx=None, sigs=None, sigp=None, zn=None, momap=None, refpts=None, **kwargs): """Touschek lifetime calculation

Computes the touschek lifetime using the Piwinski formula

args:
    ring:            ring use for tracking
    emity:           verticla emittance
    bunch_curr:      bunch current

keyword args:
    emitx=None:      horizontal emittance
    sigs=None:       rms bunch length
    sigp=None:       energy spread
    zn=None:         full ring :math:`Z/n`
    momap=None:      momentum aperture, has to be consistent with ``refpts``
                     if provided the momentum aperture is not calculated
    refpts=None:     ``refpts`` where the momentum aperture is calculated,
                     the default is to compute it for all elements in the
                     ring, ``len(refpts)>2`` is required
    resolution:      minimum distance between 2 grid points, default=1.0e-3
    amplitude:       max. amplitude for ``RADIAL`` and ``CARTESIAN`` or
                     initial step in ``RECURSIVE``
                     default = 0.1
    nturns=1024:     Number of turns for the tracking
    dp=None:         static momentum offset
    grid_mode:       ``at.GridMode.CARTESIAN/RADIAL`` track full vector
                     (default). ``at.GridMode.RECURSIVE``: recursive search
    use_mp=False:    Use python multiprocessing (``patpass``, default use
                     ``lattice_pass``). In case multi-processing is not
                     enabled ``GridMode`` is forced to
                     ``RECURSIVE`` (most efficient in single core)
    verbose=False:   Print out some inform
    epsabs, epsrel:  integral absolute and relative tolerances

Returns:
    tl: touschek lifetime, double expressed in seconds 
    ma: momentum aperture (len(refpts), 2) array
    refpts: refpts used for ma calcualtion (len(refpts), ) array
"""

epsabs = kwargs.pop('epsabs', 1.0e-16)
epsrel = kwargs.pop('epsrel', 1.0e-12)
interpolate = kwargs.pop('interpolate', True)

emitx, sigs, sigp = get_beam_sizes(ring, bunch_curr, zn=zn,
                                   emitx=emitx, sigs=sigs, sigp=sigp)
if refpts is None:
    refpts = ring.get_uint32_index(at.All, endpoint=False)
else:
    refpts = ring.get_uint32_index(refpts)
if momap is not None:
    assert len(momap) == len(refpts), \
        'Input momap and refpts have different lengths'

mask = [ring[r].Length > 0.0 for r in refpts]
if not numpy.all(mask):
    zerolength_warning = ('zero-length elements removed '
                          'from lifetime calculation')
    warnings.warn(AtWarning(zerolength_warning))
refpts = refpts[mask]

if momap is None:
    resolution = kwargs.pop('resolution', 1.0e-3)
    amplitude = kwargs.pop('amplitude', 0.1)
    kwargs.update({'refpts': refpts})
    momap, _, _ = ring.get_momentum_acceptance(resolution,
                                               amplitude, **kwargs)
else:
    momap = momap[mask]

if interpolate:
    refpts_all = numpy.array([i for i in range(refpts[0],
                                               refpts[-1]+1)
                              if ring[i].Length > 0])
    spos = numpy.squeeze(ring.get_s_pos(refpts))
    spos_all = numpy.squeeze(ring.get_s_pos(refpts_all))
    momp = numpy.interp(spos_all, spos, momap[:, 0])
    momn = numpy.interp(spos_all, spos, momap[:, 1])
    momap_all = numpy.vstack((momp, momn)).T
    ma, rp = momap_all, refpts_all
else:
    ma, rp = momap, refpts

length_all = numpy.array([e.Length for e in ring[rp]])

nc = bunch_curr/ring.revolution_frequency/qe
beta2 = ring.beta*ring.beta
gamma2 = ring.gamma*ring.gamma

emit = numpy.array([emitx, emity])
_, _, ld = ring.get_optics(refpts=rp)
bxy = ld.beta
bxy2 = bxy*bxy
axy = ld.alpha
dxy = ld.dispersion[:, [0, 2]]
dxy2 = dxy*dxy
dpxy = ld.dispersion[:, [1, 3]]
sigb = numpy.sqrt(bxy*emit)
sigb2 = sigb*sigb
sigp2 = sigp*sigp
sig = numpy.sqrt(sigb2+sigp2*dxy2)
sig2 = sig*sig
dt = dxy*axy+dpxy*bxy
dt2 = dt*dt
sigh2 = 1/(1/sigp2 + numpy.sum((dxy2+dt2)/sigb2, axis=1))

bs = bxy2/sigb2*(1-(sigh2*(dt2/sigb2).T).T)
bg2i = 1/(2*beta2*gamma2)
B1 = bg2i*numpy.sum(bs, axis=1)
B2sq = bg2i*bg2i*(numpy.diff(bs, axis=1).T**2 +
                  sigh2*sigh2*numpy.prod(bxy2*dt2, axis=1) /
                  numpy.prod(sigb2*sigb2, axis=1))
B2 = numpy.squeeze(numpy.sqrt(B2sq))

invtl = numpy.zeros(2)
val2 = numpy.zeros(len(rp)) #compute local piwinski scattering rate (https://arxiv.org/abs/physics/9903034v1)
for i in range(2):
    dpp = ma[:, i]
    um = beta2*dpp*dpp
    km = numpy.arctan(numpy.sqrt(um))

    val = numpy.zeros(len(rp)) 
    for ii in range(len(rp)):
        args = (km[ii], B1[ii], B2[ii])
        val[ii], _ = integrate.quad(int_piwinski, args[0], numpy.pi/2,
                                    args=args, epsabs=epsabs,
                                    epsrel=epsrel)

    val *= (_e_radius**2*clight*nc /
            (8*numpy.pi*(gamma2)*sigs *
             numpy.sqrt(numpy.prod(sig2, axis=1) -
             sigp2*sigp2*numpy.prod(dxy2, axis=1))) *
            2*numpy.sqrt(numpy.pi*(B1*B1-B2*B2)))

    val2 += val
    invtl[i] = sum(val*length_all.T)/sum(length_all)

tl = 1/numpy.mean(invtl)

scattering_rate = val2*nc/2    # to satisfy equation (31) using equation (41) need to multiply by nc/2 
                                                     # (division by two comes from averaging over positive and negative momentum acceptance)
return tl, momap, refpts, scattering_rate 

Thanks in advance Mihail Miceski and Barbara Dalena

swhite2401 commented 1 year ago

Hello,

Changing the interface of an existing function is discouraged because it is not backward compatible, what I can propose is to provide another function that would output the quantity you need. Would that be ok for you?

Cheers

bdalena commented 1 year ago

Hi,

Yes of course, you are the code master and it is up to you to chose your best way to modify it. We didn't want to duplicate lines of the code, but maybe the function can be split in more functions which can be called by two differents methods one for lifetime and the other one to compute the local Piwinski rate for touschek scattering.

Ciao Mihail & Barbara

lfarv commented 1 year ago

Solved in #624