From 3116b13c4d91c5ac6751aa8a96af771f662c2cf5 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Tue, 21 Oct 2025 13:30:18 -0400 Subject: [PATCH 1/4] Add continuousLinear function to teapot_base --- src/teapot/teapotbase.cc | 95 +++++++++++++++++++++++++++++++++++ src/teapot/teapotbase.hh | 2 + src/teapot/wrap_teapotbase.cc | 18 +++++++ 3 files changed, 115 insertions(+) diff --git a/src/teapot/teapotbase.cc b/src/teapot/teapotbase.cc index 9f00f23c..5ea002cf 100644 --- a/src/teapot/teapotbase.cc +++ b/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 diff --git a/src/teapot/teapotbase.hh b/src/teapot/teapotbase.hh index 5a266840..351a407a 100644 --- a/src/teapot/teapotbase.hh +++ b/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 diff --git a/src/teapot/wrap_teapotbase.cc b/src/teapot/wrap_teapotbase.cc index 48b88712..da405a16 100644 --- a/src/teapot/wrap_teapotbase.cc +++ b/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 } }; From 91e8b23ad50ec893fbbf53339e10f281921d7985 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Tue, 21 Oct 2025 13:30:52 -0400 Subject: [PATCH 2/4] Add test_continuous_focusing.py --- examples/TEAPOT/test_continuous_focusing.py | 49 +++++++++++++++++++++ 1 file changed, 49 insertions(+) create mode 100644 examples/TEAPOT/test_continuous_focusing.py diff --git a/examples/TEAPOT/test_continuous_focusing.py b/examples/TEAPOT/test_continuous_focusing.py new file mode 100644 index 00000000..c4cf2409 --- /dev/null +++ b/examples/TEAPOT/test_continuous_focusing.py @@ -0,0 +1,49 @@ +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_Lattice +from orbit.utils.consts import mass_proton + + +def test_continuous_focusing_linear(): + length = 1.0 + kq = 0.5 + + 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]) + + 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) From b4b96125254713cb1ce56f4a8516bc4a657fce3e Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Tue, 21 Oct 2025 13:45:34 -0400 Subject: [PATCH 3/4] Add ContinuousFocusingTEAPOT element --- py/orbit/teapot/__init__.py | 2 + py/orbit/teapot/teapot.py | 82 +++++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+) diff --git a/py/orbit/teapot/__init__.py b/py/orbit/teapot/__init__.py index 945450d4..2642d95f 100644 --- a/py/orbit/teapot/__init__.py +++ b/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") diff --git a/py/orbit/teapot/teapot.py b/py/orbit/teapot/teapot.py index 12cd1c34..9409101a 100644 --- a/py/orbit/teapot/teapot.py +++ b/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 From c5325fe16dc6c6226a7e30def7b731ef290aaa0a Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Tue, 21 Oct 2025 13:45:56 -0400 Subject: [PATCH 4/4] Test ContinuousFocusingTEAPOT node --- examples/TEAPOT/test_continuous_focusing.py | 58 +++++++++++++++++++-- 1 file changed, 54 insertions(+), 4 deletions(-) diff --git a/examples/TEAPOT/test_continuous_focusing.py b/examples/TEAPOT/test_continuous_focusing.py index c4cf2409..fc8ba70a 100644 --- a/examples/TEAPOT/test_continuous_focusing.py +++ b/examples/TEAPOT/test_continuous_focusing.py @@ -6,14 +6,50 @@ from orbit.core.bunch import BunchTwissAnalysis from orbit.lattice import AccLattice from orbit.lattice import AccNode -from orbit.teapot import TEAPOT_Lattice +from orbit.teapot import teapot from orbit.utils.consts import mass_proton -def test_continuous_focusing_linear(): +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) @@ -37,8 +73,22 @@ def test_continuous_focusing_linear(): bunch1.addParticle(*coords[i]) bunch2.addParticle(*coords[i]) - teapot_base.continuousLinear(bunch1, length, kq) - teapot_base.quad1(bunch2, length, kq) + 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)