Add continuous linear focusing element

examples/TEAPOT/test_continuous_focusing.py

@@ -0,0 +1,99 @@
+import math
+import random
+
+from orbit.core import teapot_base
+from orbit.core.bunch import Bunch
+from orbit.core.bunch import BunchTwissAnalysis
+from orbit.lattice import AccLattice
+from orbit.lattice import AccNode
+from orbit.teapot import teapot
+from orbit.utils.consts import mass_proton
+
+
+def test_continuous_linear_focusing():
+    coords = []
+    for i in range(100):
+        x = random.gauss(mu=0.0, sigma=0.010)
+        y = random.gauss(mu=0.0, sigma=0.010)
+        z = random.gauss(mu=0.0, sigma=1.0)
+        xp = random.gauss(mu=0.0, sigma=0.010)
+        yp = random.gauss(mu=0.0, sigma=0.010)
+        dE = random.gauss(mu=0.0, sigma=0.001)
+        coords.append([x, xp, y, yp, z, dE])
+
+    bunch_init = Bunch()
+    bunch_init.mass(mass_proton)
+    bunch_init.getSyncParticle().kinEnergy(1.000)
+
+    bunch1 = Bunch()
+    bunch2 = Bunch()
+    bunch_init.copyBunchTo(bunch1)
+    bunch_init.copyBunchTo(bunch2)
+
+    for i in range(100):
+        bunch1.addParticle(*coords[i])
+        bunch2.addParticle(*coords[i])
+
+    length = 1.0
+    kq = 0.5
+
+    teapot_base.continuousLinear(bunch1, length, kq)
+    teapot_base.quad1(bunch2, length, kq)
+
+    for i in range(bunch_init.getSize()):
+        assert bunch1.x(i) == bunch2.x(i)
+        assert bunch1.y(i) == bunch2.x(i)
+        assert bunch1.z(i) == bunch2.z(i)
+        assert bunch1.xp(i) == bunch2.xp(i)
+        assert bunch1.yp(i) == bunch2.xp(i)
+        assert bunch1.dE(i) == bunch2.dE(i)
+
+
+def test_continuous_linear_focusing_node():
+    coords = []
+    for i in range(100):
+        x = random.gauss(mu=0.0, sigma=0.010)
+        y = random.gauss(mu=0.0, sigma=0.010)
+        z = random.gauss(mu=0.0, sigma=1.0)
+        xp = random.gauss(mu=0.0, sigma=0.010)
+        yp = random.gauss(mu=0.0, sigma=0.010)
+        dE = random.gauss(mu=0.0, sigma=0.001)
+        coords.append([x, xp, y, yp, z, dE])
+
+    bunch_init = Bunch()
+    bunch_init.mass(mass_proton)
+    bunch_init.getSyncParticle().kinEnergy(1.000)
+
+    bunch1 = Bunch()
+    bunch2 = Bunch()
+    bunch_init.copyBunchTo(bunch1)
+    bunch_init.copyBunchTo(bunch2)
+
+    for i in range(100):
+        bunch1.addParticle(*coords[i])
+        bunch2.addParticle(*coords[i])
+
+    length = 1.0
+    kq = 0.5
+    nparts = 10
+
+    node1 = teapot.ContinuousFocusingTEAPOT()
+    node1.setLength(length)
+    node1.setnParts(nparts)
+    node1.setParam("kq", kq)
+
+    node2 = teapot.QuadTEAPOT()
+    node2.setLength(length)
+    node2.setnParts(nparts)
+    node2.setParam("kq", kq)
+
+    node1.trackBunch(bunch1)
+    node2.trackBunch(bunch2)
+
+    for i in range(bunch_init.getSize()):
+        assert bunch1.x(i) == bunch2.x(i)
+        assert bunch1.y(i) == bunch2.x(i)
+        assert bunch1.z(i) == bunch2.z(i)
+        assert bunch1.xp(i) == bunch2.xp(i)
+        assert bunch1.yp(i) == bunch2.xp(i)
+        assert bunch1.dE(i) == bunch2.dE(i)

py/orbit/teapot/__init__.py

@@ -16,6 +16,7 @@
 from .teapot import SolenoidTEAPOT
 from .teapot import TiltTEAPOT
 from .teapot import NodeTEAPOT
+from .teapot import ContinuousFocusingTEAPOT
 
 from .teapot import TPB
 
@@ -36,5 +37,6 @@
 __all__.append("FringeFieldTEAPOT")
 __all__.append("TiltTEAPOT")
 __all__.append("NodeTEAPOT")
+__all__.append("ContinuousFocusingTEAPOT")
 __all__.append("TPB")
 __all__.append("TEAPOT_MATRIX_Lattice")

py/orbit/teapot/teapot.py

@@ -1404,3 +1404,85 @@ def getUsage(self):
         field will be used in calculation.
         """
         return self.__usage
+
+
+class ContinuousFocusingTEAPOT(NodeTEAPOT):
+    def __init__(self, name="continuous focusing no name"):
+        NodeTEAPOT.__init__(self, name)
+
+        self.addParam("kq", 0.0)
+        self.addParam("poles", [])
+        self.addParam("kls", [])
+        self.addParam("skews", [])
+        self.setnParts(2)
+        self.waveform = None
+
+        self.getNodeTiltIN().setType("continuous focusing tilt in")
+        self.getNodeTiltOUT().setType("continuous focusing tilt out")
+
+        self.setType("continuous focusing")
+
+    def initialize(self):
+        nParts = self.getnParts()
+        if nParts < 2:
+            msg = "The Quad Combined Function TEAPOT class instance should have more than 2 parts!"
+            msg = msg + os.linesep
+            msg = msg + "Method initialize():"
+            msg = msg + os.linesep
+            msg = msg + "Name of element=" + self.getName()
+            msg = msg + os.linesep
+            msg = msg + "Type of element=" + self.getType()
+            msg = msg + os.linesep
+            msg = msg + "nParts =" + str(nParts)
+            orbitFinalize(msg)
+
+        lengthIN = (self.getLength() / (nParts - 1)) / 2.0
+        lengthOUT = (self.getLength() / (nParts - 1)) / 2.0
+        lengthStep = lengthIN + lengthOUT
+        self.setLength(lengthIN, 0)
+        self.setLength(lengthOUT, nParts - 1)
+        for i in range(nParts - 2):
+            self.setLength(lengthStep, i + 1)
+
+    def track(self, paramsDict):
+        nParts = self.getnParts()
+        index = self.getActivePartIndex()
+        length = self.getLength(index)
+
+        strength = 1.0
+        if self.waveform:
+            strength = self.waveform.getStrength()
+
+        kq = strength * self.getParam("kq")
+        poleArr = self.getParam("poles")
+        klArr = self.getParam("kls")
+        skewArr = self.getParam("skews")
+
+        bunch = paramsDict["bunch"]
+
+        useCharge = 1
+        if "useCharge" in paramsDict:
+            useCharge = paramsDict["useCharge"]
+
+        if index == 0:
+            TPB.continuousLinear(bunch, length, kq, useCharge)
+            return
+        if index > 0 and index < (nParts - 1):
+            for i in range(len(poleArr)):
+                pole = poleArr[i]
+                kl = strength * klArr[i] / (nParts - 1)
+                skew = skewArr[i]
+                TPB.multp(bunch, pole, kl, skew, useCharge)
+            TPB.continuousLinear(bunch, length, kq, useCharge)
+            return
+        if index == (nParts - 1):
+            for i in range(len(poleArr)):
+                pole = poleArr[i]
+                kl = strength * klArr[i] / (nParts - 1)
+                skew = skewArr[i]
+                TPB.multp(bunch, pole, kl, skew, useCharge)
+            TPB.continuousLinear(bunch, length, kq, useCharge)
+        return
+
+    def setWaveform(self, waveform):
+        self.waveform = waveform

src/teapot/teapotbase.cc

@@ -1566,4 +1566,99 @@ void RingRF(Bunch* bunch, double ring_length, int harmonic_numb,
     }
 }
 
+
+////////////////////////////
+// NAME
+//   continuousLinear
+//
+// DESCRIPTION
+//   Continuous axisymmetric focusing element: linear transport matrix. This
+//   is equivalent to quad1, but is focusing in both planes. This element
+//   assumes there is no energy offset from the synchronous particle.
+//
+// PARAMETERS
+//   bunch  = reference to the macro-particle bunch
+//   length = length of transport
+//   kq = quadrupole field strength [m^(-2)]
+//        kq already has information about the charge of particle.
+//
+// RETURNS
+//   Nothing
+//
+///////////////////////////////////////////////////////////////////////////
+
+void continuousLinear(Bunch* bunch, double length, double kq, int useCharge) {
+    if (kq == 0.0 || bunch->getCharge() == 0.0) {
+        drift(bunch, length);
+        return;
+    }
+
+    double kqc = abs(kq);
+    double sqrt_kq;
+    double kqlength;
+
+    double x_init;
+    double y_init;
+    double xp_init;
+    double yp_init;
+
+    double cx;
+    double sx;
+    double cy;
+    double sy;
+
+    double m11 = 0.0;
+    double m12 = 0.0;
+    double m21 = 0.0;
+    double m22 = 0.0;
+    double m33 = 0.0;
+    double m34 = 0.0;
+    double m43 = 0.0;
+    double m44 = 0.0;
+
+    SyncPart* syncPart = bunch->getSyncPart();
+
+    double v = OrbitConst::c * syncPart->getBeta();
+    if (length > 0.0) {
+        syncPart->setTime(syncPart->getTime() + length / v);
+    }
+
+    double gamma2i = 1.0 / (syncPart->getGamma() * syncPart->getGamma());
+    double dp_p_coeff = 1.0 /(syncPart->getMomentum() * syncPart->getBeta());
+
+    sqrt_kq  = pow(kqc, 0.5);
+    kqlength = sqrt_kq * length;
+    cx = cos(kqlength);
+    sx = sin(kqlength);
+    cy = cos(kqlength);
+    sy = sin(kqlength);
+    m11 = +cx;
+    m12 = +sx / sqrt_kq;
+    m21 = -sx * sqrt_kq;
+    m22 = +cx;
+    m33 = +cy;
+    m34 = +sy / sqrt_kq;
+    m43 = -sy * sqrt_kq;
+    m44 = +cy;
+
+    double dp_p;
+
+    double** arr = bunch->coordArr();
+
+    for(int i = 0; i < bunch->getSize(); i++) {
+        dp_p    = arr[i][5] * dp_p_coeff;
+        x_init  = arr[i][0];
+        xp_init = arr[i][1];
+        y_init  = arr[i][2];
+        yp_init = arr[i][3];
+
+        arr[i][0]  = x_init * m11 + xp_init * m12;
+        arr[i][1]  = x_init * m21 + xp_init * m22;
+        arr[i][2]  = y_init * m33 + yp_init * m34;
+        arr[i][3]  = y_init * m43 + yp_init * m44;
+        arr[i][4] += dp_p * gamma2i * length;
+    }
+}
+
+
 }  //end of namespace teapot_base

src/teapot/teapotbase.hh

@@ -75,6 +75,8 @@ namespace teapot_base
                      int nsteps, int useCharge);
 
     void RingRF(Bunch* bunch, double ring_length, int harmonic_numb, double voltage, double phase_s, int useCharge);
+
+    void continuousLinear(Bunch* bunch, double length, double kq, int useCharge);
 }
 
 #endif  //TEAPOT_BASE_H

src/teapot/wrap_teapotbase.cc

@@ -483,6 +483,23 @@ extern "C"
         return Py_None;
     }
 
+    // Continuous axisymmetric focusing element.
+    static PyObject* wrap_continuousLinear(PyObject *self, PyObject *args)
+    {
+        PyObject* pyBunch;
+        double length, kq;
+        int useCharge = 1;
+        if(!PyArg_ParseTuple(args, "Odd|i:continuousLinear",
+                             &pyBunch, &length, &kq, &useCharge))
+        {
+            error("teapotbase - continuousLinear - cannot parse arguments!");
+        }
+        Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) pyBunch)->cpp_obj;
+        teapot_base::continuousLinear(cpp_bunch, length, kq, useCharge);
+        Py_INCREF(Py_None);
+        return Py_None;
+    }
+
     static PyMethodDef teapotbaseMethods[] =
     {
             {"rotatexy",         wrap_rotatexy,       METH_VARARGS, "Rotates bunch around z axis "},
@@ -510,6 +527,7 @@ extern "C"
             {"soln",             wrap_soln,           METH_VARARGS, "Integration through a solenoid "},
             {"wedgebendCF",      wrap_wedgebendCF,    METH_VARARGS, "Straight bends particles through wedge for Combined Function non-SBEND "},
             {"RingRF",           wrap_RingRF,         METH_VARARGS, "Tracking particles through a simple ring RF cavity."},
+            {"continuousLinear", wrap_continuousLinear, METH_VARARGS, "Continuous axisymmetric linear focusing element."},
             { NULL, NULL }
     };