Computing the kick and invariants for nonlinear integrable optics

invariants.py

"""
Class for computing the invariants and kick for the nonlinear elliptic potential.

Use of this class should reference IPAC'15 proceeding number MOPMA029.
"""

import numpy as np
from numpy import sqrt, arccosh, arccos, pi

class Invariants:

    def __init__(self, _t, _c, _beta, _betaPrime = 0.):
        self.t = _t
        self.c = _c
        self.beta = _beta
        self.betaPrime = _betaPrime

    def computeInvariant(self, xHat, pxHat, yHat, pyHat):
        """Compute the second invariant for the integrable elliptic potential"""

        # return zero, if the magnet is turned off
        #if self.t == 0.:
        #    return 0.

        # normalize the phase space coordinates

        # Angular momentum
        ang_momentum = (xHat * pyHat - yHat * pxHat)**2
        lin_momentum = (self.c * pxHat)**2

        # elliptic coordinates
        xN = xHat / self.c
        yN = yHat / self.c

        u = ( sqrt((xN + 1)**2 + yN**2) +
              sqrt((xN - 1)**2 + yN**2) )/(2.)
        v = ( sqrt((xN + 1)**2 + yN**2) -
              sqrt((xN - 1)**2 + yN**2) )/(2.)

        # harmonic part of the potential
        f1u = self.c**2 * u**2 * (u**2-1.)
        g1v = self.c**2 * v**2 * (1.-v**2)

        # elliptic part of the potential
        f2u = -self.t * self.c**2 * u * sqrt(u**2-1.) * arccosh(u)
        g2v = -self.t * self.c**2 * v * sqrt(1.-v**2) * \
                (0.5*np.pi-arccos(v))

        # combined potentials
        fu = (0.5 * f1u - f2u)
        gv = (0.5 * g1v + g2v)

        # putting it all together
        invariant = (ang_momentum+lin_momentum) + 2.*(self.c**2)*\
                    (fu * v**2 + gv * u**2)/(u**2 - v**2)
        return invariant



    def computeHamiltonian(self, xHat, pxHat, yHat, pyHat):
        """Compute the Hamiltonian (1st invariant) for the integrable elliptic potential"""

        quadratic = 0.5 * (pxHat**2 + pyHat**2) + 0.5 * (xHat**2 + yHat**2)

        elliptic = 0.
        kfac = 1.
        if self.t != 0.:
            xN = xHat / self.c
            yN = yHat / self.c

            # Elliptic coordinates
            u = ( sqrt((xN + 1.)**2 + yN**2) +
                  sqrt((xN - 1.)**2 + yN**2) )/2.
            v = ( sqrt((xN + 1.)**2 + yN**2) -
                  sqrt((xN - 1.)**2 + yN**2) )/2.

            f2u = u * sqrt(u**2 - 1.) * arccosh(u)
            g2v = v * sqrt(1. - v**2) * (-pi/2 + arccos(v))

            kfac = self.t * self.c**2
            elliptic = (f2u + g2v) / (u**2 - v**2)

        hamiltonian = quadratic + kfac * elliptic
        return hamiltonian
        
    def computeKick(self, _x, _y):
        """
        Compute the kick
        """
        # convert everything to normalized coordinates
        # if the amplitude is zero, return zeros
        if self.t == 0.:
            return 0., 0.

        # normalize to beta
        x = _x/sqrt(self.beta)
        y = _y/sqrt(self.beta)

        # normalize to the elliptic 'focal length'
        x /= self.c
        y /= self.c

        # transform to hyperbolic coordinates
        u = (sqrt((x+1.)**2 + y**2) + sqrt((x-1.)**2 + y**2))/2.
        v = (sqrt((x+1.)**2 + y**2) - sqrt((x-1.)**2 + y**2))/2.

        dd = np.zeros(u.size)
        for idx in range (0,u.size):
            if u[idx] != 1.:
                dd[idx] = u[idx]**2 * arccosh(u[idx]) / sqrt(u[idx]**2-1.)

        # Frequently recurring quantity
        uvdiff = u**2 - v**2

        # dUdu is derivative of potential with respect to u
        dUdu = (u + arccosh(u) * sqrt(u**2-1.) + dd) / uvdiff - \
               (2*u*(u*arccosh(u)*sqrt(u**2-1.)+v*(arccos(v)-0.5*pi)*sqrt(1.-v**2)))/uvdiff**2

        # dUdv is derivative of potential with respect to v
        dUdv = (2*v*(u*arccosh(u)*sqrt(u**2-1.)+v*(arccos(v)-0.5*pi)*sqrt(1.-v**2)))/uvdiff**2 - \
               (v-(arccos(v)-0.5*pi)*sqrt(1.-v**2)+(v**2)*(arccos(v)-0.5*pi)/sqrt(1.-v**2))/uvdiff

        # Derivatives of u and v with respect to x and y
        dudx = (x+1.) / (2.*sqrt(x**2+2.*x+y**2+1.)) + \
               (x-1.) / (2.*sqrt(x**2-2.*x+y**2+1.))
        dudy = y / (2.*sqrt(y**2 + (x-1.)**2)) + \
               y / (2.*sqrt(y**2 + (x+1.)**2))

        dvdx = (x+1.) / (2.*sqrt(x**2+2.*x+y**2+1.)) - \
               (x-1.) / (2.*sqrt(x**2-2.*x+y**2+1.))
        dvdy = y / (2.*sqrt(y**2 + (x+1.)**2)) - \
               y / (2.*sqrt(y**2 + (x-1.)**2))

        kfac = self.t * self.c / np.power(self.beta,1.5)
        kick_x = kfac * (dUdu*dudx + dUdv*dvdx)
        kick_y = kfac * (dUdu*dudy + dUdv*dvdy)
        return kick_x, kick_y