From 317a502a7a52af1618bac4b25aba8c722930c10a Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Wed, 2 Sep 2026 21:26:24 -0400 Subject: [PATCH 01/13] Add getCharge() method to SyncParticle --- src/orbit/SyncPart.cc | 4 +++ src/orbit/SyncPart.hh | 5 +++ src/orbit/wrap_syncpart.cc | 61 ++++++++++++++++++++------------ tests/py/orbit/test_sync_part.py | 18 ++++++++++ 4 files changed, 65 insertions(+), 23 deletions(-) create mode 100644 tests/py/orbit/test_sync_part.py diff --git a/src/orbit/SyncPart.cc b/src/orbit/SyncPart.cc index 6494ecbf..320f9850 100644 --- a/src/orbit/SyncPart.cc +++ b/src/orbit/SyncPart.cc @@ -320,6 +320,10 @@ double SyncPart::getMass(){ return bunch->getMass(); } +double SyncPart::getCharge(){ + return bunch->getCharge(); +} + void SyncPart::readSyncPart(const char* fileName){ //for MPI diff --git a/src/orbit/SyncPart.hh b/src/orbit/SyncPart.hh index d14aa9ba..2081e4bf 100644 --- a/src/orbit/SyncPart.hh +++ b/src/orbit/SyncPart.hh @@ -68,6 +68,11 @@ class SyncPart */ double getMass(); + /** + Charge in elementary charge units + */ + double getCharge(); + /** time in seconds */ diff --git a/src/orbit/wrap_syncpart.cc b/src/orbit/wrap_syncpart.cc index 787c00a7..419ae556 100644 --- a/src/orbit/wrap_syncpart.cc +++ b/src/orbit/wrap_syncpart.cc @@ -129,6 +129,20 @@ extern "C" { return Py_BuildValue("d",val); } + // charge() - returns charge in elementary charge units + static PyObject* SyncPart_charge(PyObject *self, PyObject *args){ + pyORBIT_Object* pySyncPart = (pyORBIT_Object*) self; + int nVars = PyTuple_Size(args); + if(nVars > 0){ + error("PySyncPart - charge() - you should use bunch.charge(charge_value) instead!"); + } + double val = 0.; + SyncPart* cpp_SyncPart = (SyncPart*) pySyncPart->cpp_obj; + val = cpp_SyncPart->getCharge(); + return Py_BuildValue("d",val); + } + + //Sets or returns the momentum for the SyncPart object // the action is depended on the number of arguments // momentum() - returns momentum @@ -543,37 +557,38 @@ extern "C" { return Py_BuildValue("d",val); } - // defenition of the methods of the python SyncPart wrapper class - // they will be vailable from python level + // Definition of the methods of the python SyncPart wrapper class. + // They will be available from the Python level. static PyMethodDef SyncPartClassMethods[] = { - { "mass", SyncPart_mass ,METH_VARARGS,"Returns mass in GeV"}, - { "momentum", SyncPart_momentum ,METH_VARARGS,"Returns or sets momentum in GeV/c."}, - { "beta", SyncPart_beta ,METH_VARARGS,"Returns beta=v/c"}, - { "gamma", SyncPart_gamma ,METH_VARARGS,"Returns gamma=1/sqrt(1-(v/c)**2)"}, - { "kinEnergy", SyncPart_kinEnergy ,METH_VARARGS,"Returns or sets kinetic energy of the synchronous particle in MeV"}, - { "time", SyncPart_time ,METH_VARARGS,"Sets or returns time in sec"}, - { "x", SyncPart_x ,METH_VARARGS,"Sets or returns the x-coordinate"}, - { "y", SyncPart_y ,METH_VARARGS,"Sets or returns the y-coordinate"}, - { "z", SyncPart_z ,METH_VARARGS,"Sets or returns the z-coordinate"}, - { "px", SyncPart_px ,METH_VARARGS,"Sets or returns the x-momentum"}, - { "py", SyncPart_py ,METH_VARARGS,"Sets or returns the y-momentum"}, - { "pz", SyncPart_pz ,METH_VARARGS,"Sets or returns the z-momentum"}, - { "pVector", SyncPart_pVector ,METH_VARARGS,"Sets or returns the momentum vector as a tuple"}, - { "rVector", SyncPart_rVector ,METH_VARARGS,"Sets or returns the position vector as a tuple"}, - { "nxVector", SyncPart_nxVector ,METH_VARARGS,"Sets or returns the x-axis vector as a tuple"}, - { "nyVector", SyncPart_nyVector ,METH_VARARGS,"Returns the y-axis vector as a tuple"}, - { "energyToMomentum", SyncPart_eToP ,METH_VARARGS,"Transforms the kinetic energy to momentum"}, - { "momentumToEnergy", SyncPart_pToE ,METH_VARARGS,"Transforms the momentum to kinetic energy"}, + { "mass", SyncPart_mass, METH_VARARGS,"Returns mass in GeV"}, + { "charge", SyncPart_charge, METH_VARARGS,"Returns charge in elementary charge units"}, + { "momentum", SyncPart_momentum, METH_VARARGS,"Returns or sets momentum in GeV/c."}, + { "beta", SyncPart_beta, METH_VARARGS,"Returns beta=v/c"}, + { "gamma", SyncPart_gamma, METH_VARARGS,"Returns gamma=1/sqrt(1-(v/c)**2)"}, + { "kinEnergy", SyncPart_kinEnergy, METH_VARARGS,"Returns or sets kinetic energy of the synchronous particle in MeV"}, + { "time", SyncPart_time, METH_VARARGS,"Sets or returns time in sec"}, + { "x", SyncPart_x, METH_VARARGS,"Sets or returns the x-coordinate"}, + { "y", SyncPart_y, METH_VARARGS,"Sets or returns the y-coordinate"}, + { "z", SyncPart_z, METH_VARARGS,"Sets or returns the z-coordinate"}, + { "px", SyncPart_px, METH_VARARGS,"Sets or returns the x-momentum"}, + { "py", SyncPart_py, METH_VARARGS,"Sets or returns the y-momentum"}, + { "pz", SyncPart_pz, METH_VARARGS,"Sets or returns the z-momentum"}, + { "pVector", SyncPart_pVector, METH_VARARGS,"Sets or returns the momentum vector as a tuple"}, + { "rVector", SyncPart_rVector, METH_VARARGS,"Sets or returns the position vector as a tuple"}, + { "nxVector", SyncPart_nxVector, METH_VARARGS,"Sets or returns the x-axis vector as a tuple"}, + { "nyVector", SyncPart_nyVector, METH_VARARGS,"Returns the y-axis vector as a tuple"}, + { "energyToMomentum", SyncPart_eToP, METH_VARARGS,"Transforms the kinetic energy to momentum"}, + { "momentumToEnergy", SyncPart_pToE, METH_VARARGS,"Transforms the momentum to kinetic energy"}, {NULL} }; - // defenition of the memebers of the python SyncPart wrapper class - // they will be vailable from python level + // Definition of the members of the python SyncPart wrapper class. + // They will be available from python level. static PyMemberDef SyncPartClassMembers [] = { {NULL} }; - //new python SyncPart wrapper type definition + // New Python SyncPart wrapper type definition. static PyTypeObject pyORBIT_SyncPart_Type = { PyVarObject_HEAD_INIT(NULL, 0) "SyncParticle", /*tp_name*/ diff --git a/tests/py/orbit/test_sync_part.py b/tests/py/orbit/test_sync_part.py new file mode 100644 index 00000000..f0b0ee7b --- /dev/null +++ b/tests/py/orbit/test_sync_part.py @@ -0,0 +1,18 @@ +from orbit.core.bunch import Bunch +from orbit.core.bunch import SyncParticle + + +def make_bunch(): + bunch = Bunch() + bunch.mass(0.938272) + bunch.charge(1.0) + bunch.getSyncParticle().kinEnergy(1.0) + return bunch + + +def test_get_mass_charge(): + bunch = make_bunch() + sync_part = bunch.getSyncParticle() + + assert sync_part.mass() == bunch.mass() + assert sync_part.charge() == bunch.charge() From 4bf3f1b5a4c4f56eb6b61441a2e686bfdd551816 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Wed, 2 Sep 2026 21:34:15 -0400 Subject: [PATCH 02/13] Move transfer matrix definitions to orbit.utils.matrix --- py/orbit/envelope/envelope.py | 2 +- py/orbit/envelope/meson.build | 1 - py/orbit/envelope/track.py | 38 +++--- py/orbit/envelope/utils.py | 42 ------- py/orbit/{envelope => utils}/matrix.py | 154 +++++++++++++++---------- py/orbit/utils/meson.build | 3 +- tests/py/orbit/test_env.py | 6 +- 7 files changed, 124 insertions(+), 122 deletions(-) rename py/orbit/{envelope => utils}/matrix.py (71%) diff --git a/py/orbit/envelope/envelope.py b/py/orbit/envelope/envelope.py index 003c24ab..a837b994 100644 --- a/py/orbit/envelope/envelope.py +++ b/py/orbit/envelope/envelope.py @@ -5,8 +5,8 @@ from orbit.core.bunch import Bunch from orbit.core.bunch import BunchTwissAnalysis from orbit.core.bunch import SyncParticle +from orbit.utils.matrix import convert_matrix_zp_to_dE -from .utils import convert_matrix_zp_to_dE from .utils import gen_dist from .utils import get_classical_radius from .utils import proj_cov_matrix diff --git a/py/orbit/envelope/meson.build b/py/orbit/envelope/meson.build index 90eab844..9be34b6b 100644 --- a/py/orbit/envelope/meson.build +++ b/py/orbit/envelope/meson.build @@ -1,7 +1,6 @@ py_sources = files([ '__init__.py', 'envelope.py', - 'matrix.py', 'track.py', 'utils.py' ]) diff --git a/py/orbit/envelope/track.py b/py/orbit/envelope/track.py index 16b62b9d..e9e952c6 100644 --- a/py/orbit/envelope/track.py +++ b/py/orbit/envelope/track.py @@ -8,8 +8,8 @@ from orbit.lattice import AccLattice from orbit.teapot import BendTEAPOT from orbit.py_linac.lattice import Bend as BendLINAC +from orbit.utils.matrix import get_matrix -from .matrix import get_matrix from .envelope import Envelope @@ -55,15 +55,17 @@ def track(self, envelope: Envelope, index_start: int = 0, index_stop: int = None nodes = self.lattice.getNodes() nodes = nodes[index_start : index_stop] + sync_part = envelope.sync_part + for node_index, node in enumerate(nodes): for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: envelope.transform(matrix) for part_index in range(node.getnParts()): for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: envelope.transform(matrix) @@ -78,19 +80,19 @@ def track(self, envelope: Envelope, index_start: int = 0, index_stop: int = None else: raise ValueError - matrix = get_matrix(node, envelope=envelope, part_index=part_index) + matrix = get_matrix(node, sync_part, part_index=part_index) if matrix is not None: if matrix_sc is not None: matrix = matrix @ matrix_sc envelope.transform(matrix) for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: envelope.transform(matrix) for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: envelope.transform(matrix) @@ -135,15 +137,17 @@ def update_history(envelope: Envelope, position: float) -> None: nodes = self.lattice.getNodes() nodes = nodes[index_start : index_stop] + sync_part = envelope.sync_part + for node_index, node in enumerate(nodes): for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: envelope.transform(matrix) for part_index in range(node.getnParts()): for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: envelope.transform(matrix) @@ -158,7 +162,7 @@ def update_history(envelope: Envelope, position: float) -> None: else: raise ValueError - matrix = get_matrix(node, envelope=envelope, part_index=part_index) + matrix = get_matrix(node, sync_part, part_index=part_index) if matrix is not None: if matrix_sc is not None: matrix = matrix @ matrix_sc @@ -168,12 +172,12 @@ def update_history(envelope: Envelope, position: float) -> None: update_history(envelope, path_length) for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: envelope.transform(matrix) for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: envelope.transform(matrix) @@ -187,16 +191,18 @@ def precompute_matrices(self, envelope: Envelope, index_start: int = 0, index_st nodes = self.lattice.getNodes() nodes = nodes[index_start : index_stop] + sync_part = envelope.sync_part + self.elements = [] for node_index, node in enumerate(nodes): for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: self.elements.append((child_node, matrix)) for part_index in range(node.getnParts()): for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: self.elements.append((child_node, matrix)) @@ -205,17 +211,17 @@ def precompute_matrices(self, envelope: Envelope, index_start: int = 0, index_st if length > 0: self.elements.append(("sc", length)) - matrix = get_matrix(node, envelope=envelope, part_index=part_index) + matrix = get_matrix(node, sync_part, part_index=part_index) if matrix is not None: self.elements.append((node, matrix)) for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: self.elements.append((node, matrix)) for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, envelope=envelope) + matrix = get_matrix(child_node, sync_part) if matrix is not None: self.elements.append((node, matrix)) diff --git a/py/orbit/envelope/utils.py b/py/orbit/envelope/utils.py index 71499e31..9102794e 100644 --- a/py/orbit/envelope/utils.py +++ b/py/orbit/envelope/utils.py @@ -2,7 +2,6 @@ import numpy as np from scipy.constants import epsilon_0 -from orbit.core.bunch import SyncParticle from orbit.utils.consts import charge_electron @@ -12,47 +11,6 @@ def get_classical_radius(charge: float, mass: float) -> float: return q**2 / (4.0 * math.pi * epsilon_0 * rest_energy) -def get_dp_p_coeff(sync_part: SyncParticle) -> float: - # dE/E = (beta^2) * dp/p - # dE = (beta^2 * E) * dp/p - # dE = (beta^2 * gamma * m * c^2) * dp/p - beta = sync_part.beta() - gamma = sync_part.gamma() - rest_energy = sync_part.mass() # GeV - return 1.0 / (beta**2 * gamma * rest_energy) - - -def get_zp_coeff(sync_part: SyncParticle) -> float: - # dE/E = (beta^2) * dp/p = (beta^2) * (gamma^2) z' - # dE = (beta^2 * gamma^2 * E) * z' - # dE = (beta^2 * gamma^3 * m * c^2) * z' - beta = sync_part.beta() - gamma = sync_part.gamma() - rest_energy = sync_part.mass() - return 1.0 / (beta**2 * gamma**3 * rest_energy) - - -def convert_matrix_dp_p_to_dE(matrix: np.ndarray, sync_part: SyncParticle) -> np.ndarray: - # v = [x, x', y, y', z, dp/p] - # w = [x, x', y, y', z, dE] - # v = A w - # v -> M v - # w -> A M A^-1 - dp_p_coeff = get_dp_p_coeff(sync_part) - matrix[:5, 5] *= dp_p_coeff - matrix[5, :5] /= dp_p_coeff - matrix[5, 6] /= dp_p_coeff # driving term - return matrix - - -def convert_matrix_zp_to_dE(matrix: np.ndarray, sync_part: SyncParticle) -> np.ndarray: - zp_coeff = get_zp_coeff(sync_part) - matrix[:5, 5] *= zp_coeff - matrix[5, :5] /= zp_coeff - matrix[5, 6] /= zp_coeff # driving term - return matrix - - def gen_dist_gauss(size: int, cov_matrix: np.ndarray) -> np.ndarray: return np.random.multivariate_normal( mean=np.zeros(cov_matrix.shape[0]), diff --git a/py/orbit/envelope/matrix.py b/py/orbit/utils/matrix.py similarity index 71% rename from py/orbit/envelope/matrix.py rename to py/orbit/utils/matrix.py index ee0cebdb..885c85b4 100644 --- a/py/orbit/envelope/matrix.py +++ b/py/orbit/utils/matrix.py @@ -1,9 +1,13 @@ +"""Transfer matrix definitions.""" import math import numpy as np from orbit.core.bunch import Bunch from orbit.core.bunch import SyncParticle +from orbit.core.orbit_utils import Matrix +from orbit.core.teapot_base import MatrixGenerator +from orbit.lattice import AccLattice from orbit.lattice import AccNode from orbit.teapot import ApertureTEAPOT from orbit.teapot import DriftTEAPOT @@ -29,12 +33,8 @@ from orbit.py_linac.lattice import TiltElement as TiltLINAC from orbit.py_linac.lattice import FringeField as FringeFieldLINAC from orbit.py_linac.lattice import BaseRF_Gap as BaseRF_Gap -from orbit.py_linac.lattice import LinacApertureNode as ApertureLINAC from orbit.utils.consts import speed_of_light -from .envelope import Envelope -from .utils import get_dp_p_coeff - IGNORE_NODE_TYPES = [ NodeTEAPOT, @@ -48,6 +48,47 @@ ] +def get_dp_p_coeff(sync_part: SyncParticle) -> float: + # dE/E = (beta^2) * dp/p + # dE = (beta^2 * E) * dp/p + # dE = (beta^2 * gamma * m * c^2) * dp/p + beta = sync_part.beta() + gamma = sync_part.gamma() + rest_energy = sync_part.mass() # GeV + return 1.0 / (beta**2 * gamma * rest_energy) + + +def get_zp_coeff(sync_part: SyncParticle) -> float: + # dE/E = (beta^2) * dp/p = (beta^2) * (gamma^2) z' + # dE = (beta^2 * gamma^2 * E) * z' + # dE = (beta^2 * gamma^3 * m * c^2) * z' + beta = sync_part.beta() + gamma = sync_part.gamma() + rest_energy = sync_part.mass() + return 1.0 / (beta**2 * gamma**3 * rest_energy) + + +def convert_matrix_dp_p_to_dE(matrix: np.ndarray, sync_part: SyncParticle) -> np.ndarray: + # v = [x, x', y, y', z, dp/p] + # w = [x, x', y, y', z, dE] + # v = A w + # v -> M v + # w -> A M A^-1 + dp_p_coeff = get_dp_p_coeff(sync_part) + matrix[:5, 5] *= dp_p_coeff + matrix[5, :5] /= dp_p_coeff + matrix[5, 6] /= dp_p_coeff # driving term + return matrix + + +def convert_matrix_zp_to_dE(matrix: np.ndarray, sync_part: SyncParticle) -> np.ndarray: + zp_coeff = get_zp_coeff(sync_part) + matrix[:5, 5] *= zp_coeff + matrix[5, :5] /= zp_coeff + matrix[5, 6] /= zp_coeff # driving term + return matrix + + def get_matrix_tilt(angle: float) -> np.ndarray: cos_phi = math.cos(angle) sin_phi = math.sin(angle) @@ -68,9 +109,7 @@ def get_matrix_kick(kx: float = 0.0, ky: float = 0.0, kE: float = 0.0) -> np.nda return M -def get_matrix_drift(envelope: Envelope, length: float) -> np.ndarray: - sync_part = envelope.sync_part - +def get_matrix_drift(sync_part: SyncParticle, length: float) -> np.ndarray: M = np.identity(7) M[0, 1] = length M[2, 3] = length @@ -81,11 +120,9 @@ def get_matrix_drift(envelope: Envelope, length: float) -> np.ndarray: return M -def get_matrix_quad(envelope: Envelope, length: float, kq: float) -> np.ndarray: +def get_matrix_quad(sync_part: SyncParticle, length: float, kq: float) -> np.ndarray: if abs(kq) == 0: - return get_matrix_drift(envelope=envelope, length=length) - - sync_part = envelope.sync_part + return get_matrix_drift(sync_part, length) sqrt_abs_kq = math.sqrt(abs(kq)) @@ -124,9 +161,7 @@ def get_matrix_quad(envelope: Envelope, length: float, kq: float) -> np.ndarray: return M -def get_matrix_bend(envelope: Envelope, length: float, theta: float) -> np.ndarray: - sync_part = envelope.sync_part - +def get_matrix_bend(sync_part: SyncParticle, length: float, theta: float) -> np.ndarray: rho = length / theta cx = math.cos(theta) sx = math.sin(theta) @@ -148,11 +183,9 @@ def get_matrix_bend(envelope: Envelope, length: float, theta: float) -> np.ndarr return M -def get_matrix_solenoid(envelope: Envelope, length: float, B: float) -> np.ndarray: +def get_matrix_solenoid(sync_part: SyncParticle, length: float, B: float) -> np.ndarray: if B == 0: - return get_matrix_drift(envelope=envelope, length=length) - - sync_part = envelope.sync_part + return get_matrix_drift(sync_part, length) phase = B * length @@ -183,11 +216,9 @@ def get_matrix_solenoid(envelope: Envelope, length: float, B: float) -> np.ndarr return M -def get_matrix_cf(envelope: Envelope, length: float, kq: float) -> np.ndarray: +def get_matrix_cf(sync_part: SyncParticle, length: float, kq: float) -> np.ndarray: if kq == 0: - return get_matrix_drift(envelope=envelope, length=length) - - sync_part = envelope.sync_part + return get_matrix_drift(sync_part, length) sqrt_abs_kq = math.sqrt(abs(kq)) @@ -206,13 +237,11 @@ def get_matrix_cf(envelope: Envelope, length: float, kq: float) -> np.ndarray: return M -def get_matrix_rf_gap(envelope: Envelope, frequency: float, E0TL: float, phase: float) -> np.ndarray: - sync_part = envelope.sync_part - +def get_matrix_rf_gap(sync_part: SyncParticle, frequency: float, E0TL: float, phase: float) -> np.ndarray: gamma = sync_part.gamma() beta = sync_part.beta() mass = sync_part.mass() - charge = envelope.charge + charge = sync_part.charge() kin_energy_in = sync_part.kinEnergy() charge_E0TL_sin = charge * E0TL * math.sin(phase) @@ -250,23 +279,33 @@ def get_matrix_rf_gap(envelope: Envelope, frequency: float, E0TL: float, phase: return M -def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.ndarray | None: +def orbit_matrix_to_numpy(matrix: Matrix) -> np.ndarray: + matrix_out = np.zeros(matrix.size()) + for i in range(matrix_out.shape[0]): + for j in range(matrix_out.shape[1]): + matrix_out[i, j] = matrix.get(i, j) + return matrix_out + + +def get_matrix(node: AccNode, sync_part: SyncParticle, part_index: int = -1, fit: bool = False) -> np.ndarray | None: """Calculate transfer matrix and update synchronous particle. - This function maps various accelerator nodes to 7 x 7 transfer matrices - for envelope tracking. For non-accelerating, finite-length nodes, the - synchronous particle time is updated as in a drift. Accelerating nodes - such as RF gaps will update the synchronous particle energy. + This function maps various accelerator nodes to 7 x 7 transfer matrices. + For non-accelerating, finite-length nodes, the synchronous particle time + is updated as in a drift. Accelerating nodes such as RF gaps will update + the synchronous particle energy. Args: node: The accelerator node. - envelope: The beam envelope. + sync_part: Synchronous particle part_index: Index of the part within the node. An index of -1 returns the transfer matrix for the entire node. + fit: Whether to fit matrix to tracking data. Returns: - 7 x 7 transfer matrix or None. If None, the node can be ignored during - envelope tracking. + 7 x 7 transfer matrix or None. """ + if fit: + raise NotImplementedError node_type = type(node) if node_type in IGNORE_NODE_TYPES: @@ -278,7 +317,7 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd if node_type is DriftTEAPOT: if length <= 0: return None - return get_matrix_drift(envelope=envelope, length=length) + return get_matrix_drift(sync_part, length) elif node_type is SolenoidTEAPOT: if length <= 0: @@ -287,16 +326,16 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd B = node.getParam("B") if node.waveform: B *= node.waveform.getStrength() - B *= envelope.charge_sign + B *= np.sign(sync_part.charge()) - return get_matrix_solenoid(envelope=envelope, length=length, B=B) + return get_matrix_solenoid(sync_part, length=length, B=B) elif node_type is MultipoleTEAPOT: if length <= 0: return None if np.all(np.abs(node.getParam("kls")) == 0): - return get_matrix_drift(envelope=envelope, length=length) + return get_matrix_drift(sync_part, length) elif node_type is QuadTEAPOT: if length <= 0: @@ -305,9 +344,9 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd kq = node.getParam("kq") if node.waveform: kq *= node.waveform.getStrength() - kq *= envelope.charge_sign + kq *= np.sign(sync_part.charge()) - return get_matrix_quad(envelope=envelope, length=length, kq=kq) + return get_matrix_quad(sync_part, length=length, kq=kq) elif node_type is BendTEAPOT: if length <= 0: @@ -316,9 +355,9 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd theta = node.getParam("theta") / (nparts - 1) if part_index == 0 or part_index == nparts - 1: theta *= 0.5 - theta *= envelope.charge_sign + theta *= np.sign(sync_part.charge()) - return get_matrix_bend(envelope=envelope, length=length, theta=theta) + return get_matrix_bend(sync_part, length=length, theta=theta) elif node_type is KickTEAPOT: scale = 1.0 @@ -333,10 +372,10 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd if abs(kx) > 0 or abs(ky) > 0 or abs(kE) > 0: return np.matmul( get_matrix_kick(kx=kx, ky=ky, kE=kE), - get_matrix_drift(envelope=envelope, length=length), + get_matrix_drift(sync_part, length), ) else: - return get_matrix_drift(envelope=envelope, length=length) + return get_matrix_drift(sync_part, length) elif node_type is TiltTEAPOT: angle = node.getTiltAngle() @@ -349,24 +388,24 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd return None kq = node.getParam("kq") - kq *= envelope.charge_sign + kq *= np.sign(sync_part.charge()) if node.waveform: kq *= node.waveform.getStrength() - return get_matrix_cf(envelope=envelope, length=length, kq=kq) + return get_matrix_cf(sync_part, length=length, kq=kq) elif node_type is DriftLINAC: if length <= 0: return None - return get_matrix_drift(envelope=envelope, length=length) + return get_matrix_drift(sync_part, length) elif node_type is QuadLINAC: if length <= 0: return None - brho = 3.335640952 * envelope.momentum / envelope.charge + brho = 3.335640952 * sync_part.momentum() / sync_part.charge() kq = node.getParam("dB/dr") / brho - return get_matrix_quad(envelope=envelope, length=length, kq=kq) + return get_matrix_quad(sync_part, length=length, kq=kq) elif node_type is BendLINAC: if length <= 0: @@ -375,14 +414,14 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd theta = node.getParam("theta") / (nparts - 1) if part_index == 0 or part_index == nparts - 1: theta *= 0.5 - theta *= envelope.charge_sign + theta *= np.sign(sync_part.charge()) - return get_matrix_bend(envelope=envelope, length=length, theta=theta) + return get_matrix_bend(sync_part, length=length, theta=theta) elif node_type is DCorrectorHLINAC: length = node.getParam("effLength") / nparts field = node.getParam("B") - delta_xp = -field * envelope.charge * length * 0.299792 / envelope.momentum + delta_xp = -field * sync_part.charge() * length * 0.299792 / sync_part.momentum() if delta_xp == 0: return None return get_matrix_kick(kx=delta_xp, ky=0.0, kE=0.0) @@ -390,7 +429,7 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd elif node_type is DCorrectorVLINAC: length = node.getParam("effLength") / nparts field = node.getParam("B") - delta_yp = -field * envelope.charge * length * 0.299792 / envelope.momentum + delta_yp = -field * sync_part.charge() * length * 0.299792 / sync_part.momentum() if delta_yp == 0: return None return get_matrix_kick(kx=0.0, ky=delta_yp, kE=0.0) @@ -398,8 +437,8 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd elif node_type is SolenoidLINAC: if length <= 0: return None - B = node.getParam("B") * envelope.charge_sign - return get_matrix_solenoid(envelope=envelope, length=length, B=B) + B = node.getParam("B") * np.sign(sync_part.charge()) + return get_matrix_solenoid(sync_part, length=length, B=B) elif node_type is TiltLINAC: angle = node.getTiltAngle() @@ -416,7 +455,6 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd phase = cavity.getPhase() + mode_phase amplitude = cavity.getAmp() - sync_part = envelope.sync_part arrival_time = sync_part.time() arrival_time_design = cavity.getDesignArrivalTime() @@ -424,7 +462,7 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd if cavity.isDesignSetUp(): phase = math.fmod(frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase, 2.0 * math.pi) else: - orbitFinalize("Run `trackDesign` first to initialize cavity phases.") + raise ValueError("Run `trackDesign` first to initialize cavity phases.") else: phase = math.fmod(frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase,2.0 * math.pi) @@ -434,7 +472,7 @@ def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.nd return None return get_matrix_rf_gap( - envelope=envelope, + sync_part=sync_part, frequency=frequency, E0TL=(E0TL * amplitude), phase=phase, diff --git a/py/orbit/utils/meson.build b/py/orbit/utils/meson.build index 6d969897..7ada54d2 100644 --- a/py/orbit/utils/meson.build +++ b/py/orbit/utils/meson.build @@ -12,7 +12,8 @@ py_sources = files([ 'consts.py', 'multiDimArray.py', 'NamedObject.py', - 'orbitFinalize.py' + 'orbitFinalize.py', + 'matrix.py' ]) python.install_sources( diff --git a/tests/py/orbit/test_env.py b/tests/py/orbit/test_env.py index 5275ab6c..403293fc 100644 --- a/tests/py/orbit/test_env.py +++ b/tests/py/orbit/test_env.py @@ -357,14 +357,14 @@ def test_rf_gap_matrix( coords_out_1 = collect_bunch(bunch_out_1)["coords"] - from orbit.envelope.matrix import get_matrix_rf_gap + from orbit.utils.matrix import get_matrix_rf_gap bunch_out_2 = Bunch() bunch_in.copyBunchTo(bunch_out_2) - envelope = Envelope(bunch=bunch_in) + sync_part = bunch_in.getSyncParticle() matrix = get_matrix_rf_gap( - envelope=envelope, + sync_part, frequency=frequency, E0TL=E0TL, phase=phase, From 6511c23e81509168a6e87d2fdd29eeccd526f591 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Wed, 2 Sep 2026 21:36:32 -0400 Subject: [PATCH 03/13] Envelope stores SyncParticle, not Bunch --- py/orbit/envelope/envelope.py | 36 +++++++++-------------------------- 1 file changed, 9 insertions(+), 27 deletions(-) diff --git a/py/orbit/envelope/envelope.py b/py/orbit/envelope/envelope.py index a837b994..ab994800 100644 --- a/py/orbit/envelope/envelope.py +++ b/py/orbit/envelope/envelope.py @@ -1,3 +1,5 @@ +from multiprocessing.sharedctypes import Synchronized + import numpy as np import scipy.constants import scipy.special @@ -26,7 +28,7 @@ class Envelope: """Represents beam envelope/centroid. Attributes: - bunch: Bunch containing synchronous particle and (optionally) test particles. + sync_part: Synchronous particle. cov_matrix: 6 x 6 covariance matrix centroid: 6 x 1 centroid vector. intensity: Total number of particles. @@ -34,41 +36,21 @@ class Envelope: def __init__( self, - bunch: Bunch, + sync_part: SynchParticle, cov_matrix: np.ndarray = None, centroid: np.ndarray = None, intensity: float = 0.0, ) -> None: - empty_bunch = Bunch() - bunch.copyEmptyBunchTo(empty_bunch) - - self.bunch = empty_bunch - self.sync_part = self.bunch.getSyncParticle() + self.sync_part = sync_part self.centroid = centroid if self.centroid is None: - if bunch.getSize(): - twiss_calc = BunchTwissAnalysis() - twiss_calc.analyzeBunch(bunch) - self.centroid = np.zeros(6) - for i in range(6): - self.centroid[i] = twiss_calc.getAverage(i) - else: - self.centroid = np.zeros(6) + self.centroid = np.zeros(6) self.cov_matrix = cov_matrix if self.cov_matrix is None: - if bunch.getSize(): - twiss_calc = BunchTwissAnalysis() - twiss_calc.analyzeBunch(bunch) - self.cov_matrix = np.zeros((6, 6)) - for i in range(6): - for j in range(6): - self.cov_matrix[i, j] = twiss_calc.getCorrelation(i, j) - self.cov_matrix[j, i] = self.cov_matrix[i, j] - else: - self.cov_matrix = np.eye(6) + self.cov_matrix = np.eye(6) self.intensity = intensity self.classical_radius = get_classical_radius(self.charge, self.mass) @@ -81,7 +63,7 @@ def __init__( def copy(self): return Envelope( - bunch=self.bunch, + sync_part=self.sync_part, cov_matrix=self.cov_matrix, centroid=self.centroid, intensity=self.intensity @@ -105,7 +87,7 @@ def mass(self) -> float: @property def charge(self) -> float: - return self.bunch.charge() + return self.sync_part.charge() @property def momentum(self) -> float: From 18dbc5f78673a0b64d0800e9dc5444a7b36bbee6 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Wed, 2 Sep 2026 21:42:39 -0400 Subject: [PATCH 04/13] Envelope now takes SyncParticle instead of Bunch in constructor --- examples/Envelope/sns_linac/test_sns_linac.py | 2 +- examples/Envelope/sns_ring/test_sns_ring.py | 2 +- .../Envelope/sns_ring/test_sns_ring_speed.py | 2 +- examples/Envelope/test_env_2d_fodo.py | 2 +- examples/Envelope/test_env_2d_fodo_speed.py | 2 +- examples/Envelope/test_env_3d_drift.py | 2 +- py/orbit/envelope/envelope.py | 4 +--- tests/py/orbit/test_env.py | 18 ++++++++++++------ 8 files changed, 19 insertions(+), 15 deletions(-) diff --git a/examples/Envelope/sns_linac/test_sns_linac.py b/examples/Envelope/sns_linac/test_sns_linac.py index 7614d54f..d057a77f 100755 --- a/examples/Envelope/sns_linac/test_sns_linac.py +++ b/examples/Envelope/sns_linac/test_sns_linac.py @@ -166,7 +166,7 @@ def main(args: argparse.Namespace) -> None: for j in range(6): cov_matrix[i, j] = cov_matrix[j, i] = twiss_calc.getCorrelation(i, j) - envelope = Envelope(bunch=bunch, cov_matrix=cov_matrix, intensity=intensity) + envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix, intensity=intensity) tracker = EnvelopeTracker(lattice, sc=("3d" if args.sc else None)) diff --git a/examples/Envelope/sns_ring/test_sns_ring.py b/examples/Envelope/sns_ring/test_sns_ring.py index 4efe4289..7518c4e4 100644 --- a/examples/Envelope/sns_ring/test_sns_ring.py +++ b/examples/Envelope/sns_ring/test_sns_ring.py @@ -157,7 +157,7 @@ def main(args: argparse.Namespace) -> None: print("TRACK ENVELOPE") envelope = Envelope( - bunch=bunch, + sync_part=sync_part, cov_matrix=cov_matrix_init, centroid=centroid_init, intensity=args.intensity, diff --git a/examples/Envelope/sns_ring/test_sns_ring_speed.py b/examples/Envelope/sns_ring/test_sns_ring_speed.py index 773febc6..8ce6a63c 100644 --- a/examples/Envelope/sns_ring/test_sns_ring_speed.py +++ b/examples/Envelope/sns_ring/test_sns_ring_speed.py @@ -80,7 +80,7 @@ print("ENVELOPE") envelope = Envelope( - bunch=bunch, + sync_part=sync_part, cov_matrix=cov_matrix_init, intensity=args.intensity, ) diff --git a/examples/Envelope/test_env_2d_fodo.py b/examples/Envelope/test_env_2d_fodo.py index 824cc3b7..0a4079ac 100644 --- a/examples/Envelope/test_env_2d_fodo.py +++ b/examples/Envelope/test_env_2d_fodo.py @@ -135,7 +135,7 @@ def main(args: argparse.Namespace) -> None: # Create envelope envelope = Envelope( - bunch=bunch, + sync_part=sync_part, cov_matrix=cov_matrix_init, centroid=centroid_init, intensity=args.intensity, diff --git a/examples/Envelope/test_env_2d_fodo_speed.py b/examples/Envelope/test_env_2d_fodo_speed.py index 063ddfaf..6882c210 100644 --- a/examples/Envelope/test_env_2d_fodo_speed.py +++ b/examples/Envelope/test_env_2d_fodo_speed.py @@ -82,7 +82,7 @@ print("ENVELOPE") envelope = Envelope( - bunch=bunch, + sync_part=sync_part, cov_matrix=cov_matrix_init, intensity=args.intensity, ) diff --git a/examples/Envelope/test_env_3d_drift.py b/examples/Envelope/test_env_3d_drift.py index 798d4448..c4ec8107 100644 --- a/examples/Envelope/test_env_3d_drift.py +++ b/examples/Envelope/test_env_3d_drift.py @@ -120,7 +120,7 @@ def main(args: argparse.Namespace) -> None: centroid_init = np.zeros(6) envelope = Envelope( - bunch=bunch, + sync_part=sync_part, cov_matrix=cov_matrix_init, centroid=centroid_init, intensity=args.intensity, diff --git a/py/orbit/envelope/envelope.py b/py/orbit/envelope/envelope.py index ab994800..6c4e962c 100644 --- a/py/orbit/envelope/envelope.py +++ b/py/orbit/envelope/envelope.py @@ -4,8 +4,6 @@ import scipy.constants import scipy.special -from orbit.core.bunch import Bunch -from orbit.core.bunch import BunchTwissAnalysis from orbit.core.bunch import SyncParticle from orbit.utils.matrix import convert_matrix_zp_to_dE @@ -36,7 +34,7 @@ class Envelope: def __init__( self, - sync_part: SynchParticle, + sync_part: SyncParticle, cov_matrix: np.ndarray = None, centroid: np.ndarray = None, intensity: float = 0.0, diff --git a/tests/py/orbit/test_env.py b/tests/py/orbit/test_env.py index 403293fc..f4367cb7 100644 --- a/tests/py/orbit/test_env.py +++ b/tests/py/orbit/test_env.py @@ -88,7 +88,9 @@ def track_and_compare_rms( # Initialize bunch bunch = Bunch() bunch.mass(mass_proton) - bunch.getSyncParticle().kinEnergy(kin_energy) + + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(kin_energy) # Track bunch particles = np.random.multivariate_normal(np.zeros(6), cov_matrix, size=nparts) @@ -103,7 +105,7 @@ def track_and_compare_rms( data["bunch"]["cov"]["out"] = cov_scale * calc_bunch_cov(bunch) # Track envelope - envelope = Envelope(bunch=bunch, cov_matrix=cov_matrix) + envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix) envelope_tracker = EnvelopeTracker(lattice=lattice) data["env"]["cov"]["in"] = cov_scale * envelope.cov_matrix @@ -386,10 +388,12 @@ def test_sc_3d_cold_expansion(): def test_track_sublattice_no_error(): bunch = Bunch() bunch.mass(mass_proton) - bunch.getSyncParticle().kinEnergy(0.001) + + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(0.001) cov_matrix = np.diag(np.square([1e-3, 0, 1e-3, 0.0, 1e-3, 0.0])) - envelope = Envelope(bunch, cov_matrix=cov_matrix) + envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix) lattice = TEAPOT_Lattice() @@ -408,10 +412,12 @@ def test_get_total_matrix() -> None: bunch = Bunch() bunch.mass(mass_proton) - bunch.getSyncParticle().kinEnergy(0.001) + + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(0.001) cov_matrix = make_default_cov_matrix() - envelope = Envelope(bunch, cov_matrix=cov_matrix, intensity=1e7) + envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix, intensity=1e7) tracker = EnvelopeTracker(lattice, sc="2d") From 888b31c00f92867dc6037e68dadde5d8af86e09a Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Wed, 2 Sep 2026 22:24:57 -0400 Subject: [PATCH 05/13] Move get_matrix function to node method node.getMatrix() --- py/orbit/envelope/track.py | 31 ++- py/orbit/lattice/AccNode.py | 15 +- py/orbit/py_linac/lattice/LinacAccNodes.py | 73 ++++++ py/orbit/py_linac/lattice/LinacRfGapNodes.py | 42 ++++ py/orbit/teapot/teapot.py | 140 ++++++++++- py/orbit/utils/matrix.py | 234 ------------------- 6 files changed, 274 insertions(+), 261 deletions(-) diff --git a/py/orbit/envelope/track.py b/py/orbit/envelope/track.py index e9e952c6..0ccc4844 100644 --- a/py/orbit/envelope/track.py +++ b/py/orbit/envelope/track.py @@ -8,7 +8,6 @@ from orbit.lattice import AccLattice from orbit.teapot import BendTEAPOT from orbit.py_linac.lattice import Bend as BendLINAC -from orbit.utils.matrix import get_matrix from .envelope import Envelope @@ -59,13 +58,13 @@ def track(self, envelope: Envelope, index_start: int = 0, index_stop: int = None for node_index, node in enumerate(nodes): for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: envelope.transform(matrix) for part_index in range(node.getnParts()): for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: envelope.transform(matrix) @@ -80,19 +79,19 @@ def track(self, envelope: Envelope, index_start: int = 0, index_stop: int = None else: raise ValueError - matrix = get_matrix(node, sync_part, part_index=part_index) + matrix = node.getMatrix(sync_part, part_index=part_index) if matrix is not None: if matrix_sc is not None: matrix = matrix @ matrix_sc envelope.transform(matrix) for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: envelope.transform(matrix) for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: envelope.transform(matrix) @@ -141,13 +140,13 @@ def update_history(envelope: Envelope, position: float) -> None: for node_index, node in enumerate(nodes): for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: envelope.transform(matrix) for part_index in range(node.getnParts()): for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: envelope.transform(matrix) @@ -162,7 +161,7 @@ def update_history(envelope: Envelope, position: float) -> None: else: raise ValueError - matrix = get_matrix(node, sync_part, part_index=part_index) + matrix = node.getMatrix(sync_part, part_index=part_index) if matrix is not None: if matrix_sc is not None: matrix = matrix @ matrix_sc @@ -172,12 +171,12 @@ def update_history(envelope: Envelope, position: float) -> None: update_history(envelope, path_length) for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: envelope.transform(matrix) for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: envelope.transform(matrix) @@ -196,13 +195,13 @@ def precompute_matrices(self, envelope: Envelope, index_start: int = 0, index_st self.elements = [] for node_index, node in enumerate(nodes): for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: self.elements.append((child_node, matrix)) for part_index in range(node.getnParts()): for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: self.elements.append((child_node, matrix)) @@ -211,17 +210,17 @@ def precompute_matrices(self, envelope: Envelope, index_start: int = 0, index_st if length > 0: self.elements.append(("sc", length)) - matrix = get_matrix(node, sync_part, part_index=part_index) + matrix = node.getMatrix(sync_part, part_index=part_index) if matrix is not None: self.elements.append((node, matrix)) for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: self.elements.append((node, matrix)) for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, sync_part) + matrix = child_node.getMatrix(sync_part) if matrix is not None: self.elements.append((node, matrix)) diff --git a/py/orbit/lattice/AccNode.py b/py/orbit/lattice/AccNode.py index fdc145dc..1926ff9c 100644 --- a/py/orbit/lattice/AccNode.py +++ b/py/orbit/lattice/AccNode.py @@ -1,13 +1,20 @@ +from __future__ import annotations + import sys import os import math +from typing import TYPE_CHECKING from ..utils import orbitFinalize from ..utils import NamedObject from ..utils import TypedObject from ..utils import ParamsDictObject -from ..lattice import AccActionsContainer +from .AccActionsContainer import AccActionsContainer + +if TYPE_CHECKING: + import numpy as np + from orbit.core.bunch import SyncParticle class AccNode(NamedObject, TypedObject, ParamsDictObject): @@ -145,6 +152,12 @@ def initialize(self): """ pass + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + """ + Return the transfer matrix for this node and update the synchronous particle. + """ + raise NotImplementedError(str(self)) + def getNumberOfChildren(self): """ Returns the total number of direct children diff --git a/py/orbit/py_linac/lattice/LinacAccNodes.py b/py/orbit/py_linac/lattice/LinacAccNodes.py index ba72fef5..4186353f 100755 --- a/py/orbit/py_linac/lattice/LinacAccNodes.py +++ b/py/orbit/py_linac/lattice/LinacAccNodes.py @@ -8,16 +8,26 @@ import os import math +import numpy as np # import the finalization function from orbit.utils import orbitFinalize # import general accelerator elements and lattice from orbit.lattice import AccNode, AccActionsContainer, AccNodeBunchTracker +from orbit.core.bunch import SyncParticle # import teapot base functions from wrapper around C++ functions from orbit.teapot_base import TPB +from orbit.utils.matrix import get_matrix_bend +from orbit.utils.matrix import get_matrix_cf +from orbit.utils.matrix import get_matrix_drift +from orbit.utils.matrix import get_matrix_kick +from orbit.utils.matrix import get_matrix_quad +from orbit.utils.matrix import get_matrix_solenoid +from orbit.utils.matrix import get_matrix_tilt + # Import the linac specific tracking from linac_tracking. This module has # the following functions duplicated the original TEAPOT functions # drift - linac drift tracking @@ -137,6 +147,9 @@ def __init__(self, name="none"): BaseLinacNode.__init__(self, name) self.setType("markerLinacNode") + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + class LinacNode(BaseLinacNode): """ @@ -304,6 +317,12 @@ def track(self, paramsDict): bunch = paramsDict["bunch"] self.tracking_module.drift(bunch, length) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + return get_matrix_drift(sync_part, length) + class Quad(LinacMagnetNode): """ @@ -513,6 +532,15 @@ def getTotalField(self, z): G = self.getParam("dB/dr") return G + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + brho = 3.335640952 * sync_part.momentum() / sync_part.charge() + kq = self.getParam("dB/dr") / brho + return get_matrix_quad(sync_part, length=length, kq=kq) + class Bend(LinacMagnetNode): """ @@ -705,6 +733,19 @@ def track(self, paramsDict): TPB.bend1(bunch, length, theta / 2.0) return + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + nparts = self.getnParts() + theta = self.getParam("theta") / (nparts - 1) + if part_index == 0 or part_index == nparts - 1: + theta *= 0.5 + theta *= np.sign(sync_part.charge()) + + return get_matrix_bend(sync_part, length=length, theta=theta) + class DCorrectorH(LinacMagnetNode): """ @@ -750,6 +791,14 @@ def track(self, paramsDict): kick = -field * charge * length * 0.299792 / momentum self.tracking_module.kick(bunch, kick, 0.0, 0.0) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getParam("effLength") / self.getnParts() + field = self.getParam("B") + delta_xp = -field * sync_part.charge() * length * 0.299792 / sync_part.momentum() + if delta_xp == 0: + return None + return get_matrix_kick(kx=delta_xp, ky=0.0, kE=0.0) + class DCorrectorV(LinacMagnetNode): """ @@ -795,6 +844,14 @@ def track(self, paramsDict): kick = field * charge * length * 0.299792 / momentum self.tracking_module.kick(bunch, 0, kick, 0.0) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getParam("effLength") / self.getnParts() + field = self.getParam("B") + delta_yp = -field * sync_part.charge() * length * 0.299792 / sync_part.momentum() + if delta_yp == 0: + return None + return get_matrix_kick(kx=0.0, ky=delta_yp, kE=0.0) + class ThickKick(LinacMagnetNode): """ @@ -892,6 +949,13 @@ def track(self, paramsDict): useCharge = paramsDict["useCharge"] TPB.soln(bunch, length, B, useCharge) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + B = self.getParam("B") * np.sign(sync_part.charge()) + return get_matrix_solenoid(sync_part, length=length, B=B) + class AbstractRF_Gap(BaseLinacNode): """ This is an abstarct class for all RF Gap classes. @@ -1010,6 +1074,12 @@ def track(self, paramsDict): bunch = paramsDict["bunch"] TPB.rotatexy(bunch, self.__angle) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + angle = self.getTiltAngle() + if angle == 0: + return None + return get_matrix_tilt(angle=angle) + class FringeField(BaseLinacNode): """ @@ -1034,6 +1104,9 @@ def track(self, paramsDict): if self.__trackFunc != None and self.__usage == True: self.__trackFunc(self, paramsDict) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + def setFringeFieldFunction(self, trackFunction=None): """ Sets the fringe field function that will track the bunch through the fringe. diff --git a/py/orbit/py_linac/lattice/LinacRfGapNodes.py b/py/orbit/py_linac/lattice/LinacRfGapNodes.py index 30b1fbe8..6f482f39 100644 --- a/py/orbit/py_linac/lattice/LinacRfGapNodes.py +++ b/py/orbit/py_linac/lattice/LinacRfGapNodes.py @@ -5,6 +5,7 @@ import os import math +import numpy as np # ---- MPI module function and classes from orbit.core.orbit_mpi import mpi_comm, mpi_datatype, MPI_Comm_rank, MPI_Bcast @@ -26,6 +27,8 @@ # The abstract RF gap import from orbit.py_linac.lattice.LinacAccNodes import AbstractRF_Gap +from orbit.utils.matrix import get_matrix_rf_gap + # import teapot base functions from wrapper around C++ functions # Import the linac specific tracking from linac_tracking. This module has @@ -35,6 +38,7 @@ # quad2 - linac quad non-linear part of tracking from orbit.core.bunch import Bunch +from orbit.core.bunch import SyncParticle class BaseRF_Gap(AbstractRF_Gap): @@ -251,6 +255,44 @@ def track(self, paramsDict): # print "debug delta_time in deg=",frequency*(arrival_time - designArrivalTime)*380. # print "debug RF =",self.getName()," E0TL=",E0TL," phase=",(phase*180./math.pi - 180.)," eKin[MeV]=",bunch.getSyncParticle().kinEnergy()*1.0e+3 + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + E0TL = self.getParam("E0TL") + mode_phase = self.getParam("mode") * math.pi + + cavity = self.getRF_Cavity() + frequency = cavity.getFrequency() + phase = cavity.getPhase() + mode_phase + amplitude = cavity.getAmp() + + arrival_time = sync_part.time() + arrival_time_design = cavity.getDesignArrivalTime() + + if self.isFirstRFGap(): + if cavity.isDesignSetUp(): + phase = math.fmod( + frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase, + 2.0 * math.pi, + ) + else: + raise ValueError("Run `trackDesign` first to initialize cavity phases.") + else: + phase = math.fmod( + frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase, + 2.0 * math.pi, + ) + + self.setGapPhase(phase) + + if amplitude == 0.0: + return None + + return get_matrix_rf_gap( + sync_part=sync_part, + frequency=frequency, + E0TL=(E0TL * amplitude), + phase=phase, + ) + def trackDesign(self, paramsDict): """ The RF First Gap node setups the design time of passage diff --git a/py/orbit/teapot/teapot.py b/py/orbit/teapot/teapot.py index 63a09b27..1d64ec67 100644 --- a/py/orbit/teapot/teapot.py +++ b/py/orbit/teapot/teapot.py @@ -19,23 +19,32 @@ import sys import os import math +import numpy as np from typing import Any from typing import Callable from typing import Union -from ..lattice import AccLattice -from ..lattice import AccNode -from ..lattice import AccActionsContainer -from ..lattice import AccNodeBunchTracker -from ..teapot_base import TPB -from ..utils import orbitFinalize -from ..parsers.mad_parser import MAD_Parser -from ..parsers.mad_parser import MAD_LattElement -from ..parsers.madx_parser import MADX_Parser -from ..parsers.madx_parser import MADX_LattElement +from orbit.lattice import AccLattice +from orbit.lattice import AccNode +from orbit.lattice import AccActionsContainer +from orbit.lattice import AccNodeBunchTracker +from orbit.teapot_base import TPB +from orbit.utils import orbitFinalize +from orbit.parsers.mad_parser import MAD_Parser +from orbit.parsers.mad_parser import MAD_LattElement +from orbit.parsers.madx_parser import MADX_Parser +from orbit.parsers.madx_parser import MADX_LattElement +from orbit.utils.matrix import get_matrix_bend +from orbit.utils.matrix import get_matrix_cf +from orbit.utils.matrix import get_matrix_drift +from orbit.utils.matrix import get_matrix_kick +from orbit.utils.matrix import get_matrix_quad +from orbit.utils.matrix import get_matrix_solenoid +from orbit.utils.matrix import get_matrix_tilt from orbit.core.aperture import Aperture from orbit.core.bunch import Bunch +from orbit.core.bunch import SyncParticle from orbit.core.bunch import BunchTwissAnalysis @@ -440,6 +449,9 @@ def track(self, paramsDict: dict) -> None: turn = bunch.bunchAttrInt("TurnNumber") bunch.bunchAttrInt("TurnNumber", turn + 1) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + class NodeTEAPOT(BaseTEAPOT): def __init__(self, name: str = "no name") -> None: @@ -464,6 +476,11 @@ def __init__(self, name: str = "no name") -> None: self.addParam("tilt", self.__tiltNodeIN.getTiltAngle()) self.setType("node teapot") + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + if type(self) is NodeTEAPOT: + return None + return super().getMatrix(sync_part, part_index=part_index) + def setTiltAngle(self, angle: float = 0.0) -> None: """ Sets the tilt angle for the tilt operation. @@ -582,6 +599,12 @@ def track(self, paramsDict: dict) -> None: bunch = paramsDict["bunch"] TPB.drift(bunch, length) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + return get_matrix_drift(sync_part, length) + class ApertureTEAPOT(NodeTEAPOT): """ @@ -622,6 +645,9 @@ def track(self, paramsDict: dict) -> None: lostbunch = paramsDict["lostbunch"] self.aperture.checkBunch(bunch, lostbunch) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + class MonitorTEAPOT(NodeTEAPOT): """ @@ -649,6 +675,9 @@ def track(self, paramsDict: dict) -> None: self.addParam("yAvg", self.twiss.getAverage(2)) self.addParam("ypAvg", self.twiss.getAverage(3)) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + class BunchWrapTEAPOT(NodeTEAPOT): """ @@ -673,6 +702,9 @@ def track(self, paramsDict: dict) -> None: length = self.getParam("ring_length") TPB.wrapbunch(bunch, length) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + class SolenoidTEAPOT(NodeTEAPOT): """ @@ -721,6 +753,18 @@ def setWaveform(self, waveform: Any) -> None: """ self.waveform = waveform + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + B = self.getParam("B") + if self.waveform: + B *= self.waveform.getStrength() + B *= np.sign(sync_part.charge()) + + return get_matrix_solenoid(sync_part, length=length, B=B) + class MultipoleTEAPOT(NodeTEAPOT): """ @@ -871,6 +915,16 @@ def setWaveform(self, waveform: Any) -> None: """ self.waveform = waveform + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + if np.all(np.abs(self.getParam("kls")) == 0): + return get_matrix_drift(sync_part, length) + + raise NotImplementedError(str(self)) + class QuadTEAPOT(NodeTEAPOT): """ @@ -1034,6 +1088,18 @@ def setWaveform(self, waveform: Any) -> None: """ self.waveform = waveform + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + kq = self.getParam("kq") + if self.waveform: + kq *= self.waveform.getStrength() + kq *= np.sign(sync_part.charge()) + + return get_matrix_quad(sync_part, length=length, kq=kq) + class BendTEAPOT(NodeTEAPOT): """ @@ -1248,6 +1314,19 @@ def track(self, paramsDict: dict) -> None: TPB.bend1(bunch, length, theta / 2.0) return + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + nparts = self.getnParts() + theta = self.getParam("theta") / (nparts - 1) + if part_index == 0 or part_index == nparts - 1: + theta *= 0.5 + theta *= np.sign(sync_part.charge()) + + return get_matrix_bend(sync_part, length=length, theta=theta) + class RingRFTEAPOT(NodeTEAPOT): """ @@ -1414,6 +1493,26 @@ def setWaveform(self, waveform): """ self.waveform = waveform + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + nparts = self.getnParts() + + scale = 1.0 + if self.waveform is not None: + scale = self.waveform.getStrength() + + scale /= nparts - 1 + kx = scale * self.getParam("kx") + ky = scale * self.getParam("ky") + kE = self.getParam("dE") + + if abs(kx) > 0 or abs(ky) > 0 or abs(kE) > 0: + return np.matmul( + get_matrix_kick(kx=kx, ky=ky, kE=kE), + get_matrix_drift(sync_part, length), + ) + return get_matrix_drift(sync_part, length) + class TiltTEAPOT(BaseTEAPOT): """ @@ -1449,6 +1548,12 @@ def track(self, paramsDict: dict) -> None: bunch = paramsDict["bunch"] TPB.rotatexy(bunch, self.__angle) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + angle = self.getTiltAngle() + if angle == 0: + return None + return get_matrix_tilt(angle) + class FringeFieldTEAPOT(BaseTEAPOT): """ @@ -1478,6 +1583,9 @@ def track(self, paramsDict: dict) -> None: if self.__trackFunc != None: self.__trackFunc(self, paramsDict) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + def setFringeFieldFunction(self, trackFunction: Callable) -> None: """ Sets the fringe field function that will track the bunch through the fringe. @@ -1595,3 +1703,15 @@ def track(self, paramsDict): def setWaveform(self, waveform): self.waveform = waveform + + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + kq = self.getParam("kq") + kq *= np.sign(sync_part.charge()) + if self.waveform: + kq *= self.waveform.getStrength() + + return get_matrix_cf(sync_part, length=length, kq=kq) diff --git a/py/orbit/utils/matrix.py b/py/orbit/utils/matrix.py index 885c85b4..0f884f2f 100644 --- a/py/orbit/utils/matrix.py +++ b/py/orbit/utils/matrix.py @@ -3,51 +3,11 @@ import numpy as np -from orbit.core.bunch import Bunch from orbit.core.bunch import SyncParticle from orbit.core.orbit_utils import Matrix -from orbit.core.teapot_base import MatrixGenerator -from orbit.lattice import AccLattice -from orbit.lattice import AccNode -from orbit.teapot import ApertureTEAPOT -from orbit.teapot import DriftTEAPOT -from orbit.teapot import BendTEAPOT -from orbit.teapot import KickTEAPOT -from orbit.teapot import MonitorTEAPOT -from orbit.teapot import MultipoleTEAPOT -from orbit.teapot import NodeTEAPOT -from orbit.teapot import QuadTEAPOT -from orbit.teapot import SolenoidTEAPOT -from orbit.teapot import FringeFieldTEAPOT -from orbit.teapot import BunchWrapTEAPOT -from orbit.teapot import TiltTEAPOT -from orbit.teapot import ContinuousLinearFocusingTEAPOT -from orbit.teapot import TurnCounterTEAPOT -from orbit.py_linac.lattice import MarkerLinacNode as MarkerLINAC -from orbit.py_linac.lattice import Drift as DriftLINAC -from orbit.py_linac.lattice import Quad as QuadLINAC -from orbit.py_linac.lattice import Bend as BendLINAC -from orbit.py_linac.lattice import DCorrectorH as DCorrectorHLINAC -from orbit.py_linac.lattice import DCorrectorV as DCorrectorVLINAC -from orbit.py_linac.lattice import Solenoid as SolenoidLINAC -from orbit.py_linac.lattice import TiltElement as TiltLINAC -from orbit.py_linac.lattice import FringeField as FringeFieldLINAC -from orbit.py_linac.lattice import BaseRF_Gap as BaseRF_Gap from orbit.utils.consts import speed_of_light -IGNORE_NODE_TYPES = [ - NodeTEAPOT, - MonitorTEAPOT, - FringeFieldTEAPOT, - ApertureTEAPOT, - BunchWrapTEAPOT, - TurnCounterTEAPOT, - MarkerLINAC, - FringeFieldLINAC, -] - - def get_dp_p_coeff(sync_part: SyncParticle) -> float: # dE/E = (beta^2) * dp/p # dE = (beta^2 * E) * dp/p @@ -285,197 +245,3 @@ def orbit_matrix_to_numpy(matrix: Matrix) -> np.ndarray: for j in range(matrix_out.shape[1]): matrix_out[i, j] = matrix.get(i, j) return matrix_out - - -def get_matrix(node: AccNode, sync_part: SyncParticle, part_index: int = -1, fit: bool = False) -> np.ndarray | None: - """Calculate transfer matrix and update synchronous particle. - - This function maps various accelerator nodes to 7 x 7 transfer matrices. - For non-accelerating, finite-length nodes, the synchronous particle time - is updated as in a drift. Accelerating nodes such as RF gaps will update - the synchronous particle energy. - - Args: - node: The accelerator node. - sync_part: Synchronous particle - part_index: Index of the part within the node. An index of -1 returns - the transfer matrix for the entire node. - fit: Whether to fit matrix to tracking data. - Returns: - 7 x 7 transfer matrix or None. - """ - if fit: - raise NotImplementedError - - node_type = type(node) - if node_type in IGNORE_NODE_TYPES: - return None - - length = node.getLength(part_index) - nparts = node.getnParts() - - if node_type is DriftTEAPOT: - if length <= 0: - return None - return get_matrix_drift(sync_part, length) - - elif node_type is SolenoidTEAPOT: - if length <= 0: - return None - - B = node.getParam("B") - if node.waveform: - B *= node.waveform.getStrength() - B *= np.sign(sync_part.charge()) - - return get_matrix_solenoid(sync_part, length=length, B=B) - - elif node_type is MultipoleTEAPOT: - if length <= 0: - return None - - if np.all(np.abs(node.getParam("kls")) == 0): - return get_matrix_drift(sync_part, length) - - elif node_type is QuadTEAPOT: - if length <= 0: - return None - - kq = node.getParam("kq") - if node.waveform: - kq *= node.waveform.getStrength() - kq *= np.sign(sync_part.charge()) - - return get_matrix_quad(sync_part, length=length, kq=kq) - - elif node_type is BendTEAPOT: - if length <= 0: - return None - - theta = node.getParam("theta") / (nparts - 1) - if part_index == 0 or part_index == nparts - 1: - theta *= 0.5 - theta *= np.sign(sync_part.charge()) - - return get_matrix_bend(sync_part, length=length, theta=theta) - - elif node_type is KickTEAPOT: - scale = 1.0 - if node.waveform is not None: - scale = node.waveform.getStrength() - - scale /= (nparts - 1) - kx = scale * node.getParam("kx") - ky = scale * node.getParam("ky") - kE = node.getParam("dE") - - if abs(kx) > 0 or abs(ky) > 0 or abs(kE) > 0: - return np.matmul( - get_matrix_kick(kx=kx, ky=ky, kE=kE), - get_matrix_drift(sync_part, length), - ) - else: - return get_matrix_drift(sync_part, length) - - elif node_type is TiltTEAPOT: - angle = node.getTiltAngle() - if angle == 0: - return None - return get_matrix_tilt(angle) - - elif node_type is ContinuousLinearFocusingTEAPOT: - if length <= 0: - return None - - kq = node.getParam("kq") - kq *= np.sign(sync_part.charge()) - if node.waveform: - kq *= node.waveform.getStrength() - - return get_matrix_cf(sync_part, length=length, kq=kq) - - elif node_type is DriftLINAC: - if length <= 0: - return None - return get_matrix_drift(sync_part, length) - - elif node_type is QuadLINAC: - if length <= 0: - return None - - brho = 3.335640952 * sync_part.momentum() / sync_part.charge() - kq = node.getParam("dB/dr") / brho - return get_matrix_quad(sync_part, length=length, kq=kq) - - elif node_type is BendLINAC: - if length <= 0: - return None - - theta = node.getParam("theta") / (nparts - 1) - if part_index == 0 or part_index == nparts - 1: - theta *= 0.5 - theta *= np.sign(sync_part.charge()) - - return get_matrix_bend(sync_part, length=length, theta=theta) - - elif node_type is DCorrectorHLINAC: - length = node.getParam("effLength") / nparts - field = node.getParam("B") - delta_xp = -field * sync_part.charge() * length * 0.299792 / sync_part.momentum() - if delta_xp == 0: - return None - return get_matrix_kick(kx=delta_xp, ky=0.0, kE=0.0) - - elif node_type is DCorrectorVLINAC: - length = node.getParam("effLength") / nparts - field = node.getParam("B") - delta_yp = -field * sync_part.charge() * length * 0.299792 / sync_part.momentum() - if delta_yp == 0: - return None - return get_matrix_kick(kx=0.0, ky=delta_yp, kE=0.0) - - elif node_type is SolenoidLINAC: - if length <= 0: - return None - B = node.getParam("B") * np.sign(sync_part.charge()) - return get_matrix_solenoid(sync_part, length=length, B=B) - - elif node_type is TiltLINAC: - angle = node.getTiltAngle() - if angle == 0: - return None - return get_matrix_tilt(angle=angle) - - elif node_type is BaseRF_Gap: - E0TL = node.getParam("E0TL") - mode_phase = node.getParam("mode") * math.pi - - cavity = node.getRF_Cavity() - frequency = cavity.getFrequency() - phase = cavity.getPhase() + mode_phase - amplitude = cavity.getAmp() - - arrival_time = sync_part.time() - arrival_time_design = cavity.getDesignArrivalTime() - - if node.isFirstRFGap(): - if cavity.isDesignSetUp(): - phase = math.fmod(frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase, 2.0 * math.pi) - else: - raise ValueError("Run `trackDesign` first to initialize cavity phases.") - else: - phase = math.fmod(frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase,2.0 * math.pi) - - node.setGapPhase(phase) - - if amplitude == 0.0: - return None - - return get_matrix_rf_gap( - sync_part=sync_part, - frequency=frequency, - E0TL=(E0TL * amplitude), - phase=phase, - ) - - raise NotImplementedError(str(node)) \ No newline at end of file From ac1686a1ecfa1ccd04269c94aba6ea3fe248ff01 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Wed, 2 Sep 2026 23:53:34 -0400 Subject: [PATCH 06/13] Track using lattice.trackEnvelope --- examples/Envelope/run.sh | 25 +- examples/Envelope/sns_linac/test_sns_linac.py | 140 +++++---- examples/Envelope/sns_ring/test_sns_ring.py | 11 +- .../Envelope/sns_ring/test_sns_ring_speed.py | 11 +- examples/Envelope/test_env_2d_fodo.py | 14 +- examples/Envelope/test_env_2d_fodo_speed.py | 5 +- examples/Envelope/test_env_3d_drift.py | 16 +- py/orbit/envelope/__init__.py | 1 - py/orbit/envelope/envelope.py | 1 + py/orbit/envelope/meson.build | 1 - py/orbit/envelope/track.py | 295 ------------------ py/orbit/lattice/AccLattice.py | 287 ++++++++++++++++- tests/py/orbit/test_env.py | 17 +- 13 files changed, 421 insertions(+), 403 deletions(-) delete mode 100644 py/orbit/envelope/track.py diff --git a/examples/Envelope/run.sh b/examples/Envelope/run.sh index 98cf808c..31054736 100755 --- a/examples/Envelope/run.sh +++ b/examples/Envelope/run.sh @@ -2,15 +2,28 @@ set -x -python test_env_2d_fodo.py -python test_env_2d_fodo_speed.py -python test_env_3d_drift.py +python test_env_2d_fodo.py --sc 0 +python test_env_2d_fodo.py --sc 1 +python test_env_2d_fodo.py --sc 1 --offset-x 0.001 +python test_env_2d_fodo.py --sc 1 --tilt 45.0 +python test_env_2d_fodo_speed.py --sc 0 +python test_env_2d_fodo_speed.py --sc 1 +python test_env_3d_drift.py --sc 0 +python test_env_3d_drift.py --sc 1 +python test_env_3d_drift.py --sc 1 --rms-y 0.002 --tilt-z 45.0 +python test_env_3d_drift.py --sc 1 --rms-z 0.002 --tilt-x 45.0 cd sns_linac -python test_sns_linac.py +python test_sns_linac.py --sc 0 +python test_sns_linac.py --sc 1 --dist gauss +python test_sns_linac.py --sc 1 --dist waterbag +python test_sns_linac.py --sc 1 --dist kv cd .. cd sns_ring -python test_sns_ring.py -python test_sns_ring_speed.py +python test_sns_ring.py --sc 0 +python test_sns_ring.py --sc 1 +python test_sns_ring.py --sc 1 --tilt 45.0 +python test_sns_ring_speed.py --sc 0 +python test_sns_ring_speed.py --sc 1 cd .. diff --git a/examples/Envelope/sns_linac/test_sns_linac.py b/examples/Envelope/sns_linac/test_sns_linac.py index d057a77f..ee644659 100755 --- a/examples/Envelope/sns_linac/test_sns_linac.py +++ b/examples/Envelope/sns_linac/test_sns_linac.py @@ -11,8 +11,10 @@ import argparse import math import os +import pathlib import random import sys +import time import numpy as np import matplotlib.pyplot as plt @@ -30,7 +32,6 @@ from orbit.bunch_generators import KVDist3D from orbit.bunch_utils import collect_bunch from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.lattice import AccLattice from orbit.lattice import AccNode from orbit.lattice import AccActionsContainer @@ -66,16 +67,47 @@ def parse_args() -> argparse.Namespace: return parser.parse_args() -def main(args: argparse.Namespace) -> None: +def make_lattice(args: argparse.Namespace) -> LinacAccLattice: + seq_names = [ + "MEBT", + "DTL1", + "DTL2", + "DTL3", + "DTL4", + "DTL5", + "DTL6", + "CCL1", + "CCL2", + "CCL3", + "CCL4", + "SCLMed", + "SCLHigh", + "HEBT1", + "HEBT2", + ] + if args.seq_stop: + index = seq_names.index(args.seq_stop) + 1 + seq_names = seq_names[:index] - output_dir = "outputs" - os.makedirs(output_dir, exist_ok=True) + sns_linac_factory = SNS_LinacLatticeFactory() + sns_linac_factory.setMaxDriftLength(args.sc_path_length_min) + lattice = sns_linac_factory.getLinacAccLattice(seq_names, "inputs/sns_linac.xml") - random.seed(23) + for node in lattice.getNodes(): + try: + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + except: + pass + + rf_gaps = lattice.getRF_Gaps() + for rf_gap in rf_gaps: + rf_gap.setCppGapModel(MatrixRfGap()) + + return lattice - # Bunch - # -------------------------------------------------------------------------------- +def make_bunch(args: argparse.Namespace) -> Bunch: kin_energy = 0.0025 # [GeV] mass = mass_proton + 2.0 * mass_electron frequency = 402.5e06 @@ -110,72 +142,64 @@ def main(args: argparse.Namespace) -> None: for _ in range(args.nparts): bunch.addParticle(*dist.getCoordinates()) + return bunch - # Lattice - # -------------------------------------------------------------------------------- - seq_names = [ - "MEBT", - "DTL1", - "DTL2", - "DTL3", - "DTL4", - "DTL5", - "DTL6", - "CCL1", - "CCL2", - "CCL3", - "CCL4", - "SCLMed", - "SCLHigh", - "HEBT1", - "HEBT2", - ] - if args.seq_stop: - index = seq_names.index(args.seq_stop) + 1 - seq_names = seq_names[:index] +def get_bunch_cov_matrix(bunch: Bunch) -> Bunch: + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) - sns_linac_factory = SNS_LinacLatticeFactory() - sns_linac_factory.setMaxDriftLength(args.sc_path_length_min) - lattice = sns_linac_factory.getLinacAccLattice(seq_names, "inputs/sns_linac.xml") + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(6): + cov_matrix[i, j] = cov_matrix[j, i] = twiss_calc.getCorrelation(i, j) + return cov_matrix - for node in lattice.getNodes(): - try: - node.setUsageFringeFieldIN(False) - node.setUsageFringeFieldOUT(False) - except: - pass - rf_gaps = lattice.getRF_Gaps() - for rf_gap in rf_gaps: - rf_gap.setCppGapModel(MatrixRfGap()) +def get_bunch_centroid(bunch: Bunch) -> Bunch: + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) - for index, node in enumerate(lattice.getNodes()): - print(index, type(node), node.getName()) + centroid = np.zeros(6) + for i in range(6): + centroid[i] = twiss_calc.getAverage(i) + return centroid - lattice.trackDesignBunch(bunch) - # Track envelope - # -------------------------------------------------------------------------------- +def make_envelope(bunch: Bunch) -> Envelope: + sync_part = bunch.getSyncParticle() + cov_matrix = get_bunch_cov_matrix(bunch) + centroid = get_bunch_centroid(bunch) + intensity = bunch.getSize() + return Envelope( + sync_part, + cov_matrix=cov_matrix, + centroid=centroid, + intensity=intensity + ) - twiss_calc = BunchTwissAnalysis() - twiss_calc.analyzeBunch(bunch) - cov_matrix = np.zeros((6, 6)) - for i in range(6): - for j in range(6): - cov_matrix[i, j] = cov_matrix[j, i] = twiss_calc.getCorrelation(i, j) +def main(args: argparse.Namespace) -> None: + + path = pathlib.Path(__file__) + output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S")) + os.makedirs(output_dir, exist_ok=True) - envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix, intensity=intensity) + random.seed(23) + + # Track envelope + bunch = make_bunch(args) + envelope = make_envelope(bunch) - tracker = EnvelopeTracker(lattice, sc=("3d" if args.sc else None)) + lattice = make_lattice(args) + lattice.trackDesignBunch(bunch) histories = {} - histories["envelope"] = tracker.track_history(envelope) + histories["envelope"] = lattice.trackEnvelopeHistory(envelope) # Track bunch - # -------------------------------------------------------------------------------- - + bunch = make_bunch(args) + lattice = make_lattice(args) lattice.trackDesignBunch(bunch) if args.sc: @@ -293,4 +317,4 @@ def main(args: argparse.Namespace) -> None: if __name__ == "__main__": - main(parse_args()) + main(parse_args()) \ No newline at end of file diff --git a/examples/Envelope/sns_ring/test_sns_ring.py b/examples/Envelope/sns_ring/test_sns_ring.py index 7518c4e4..3fafa636 100644 --- a/examples/Envelope/sns_ring/test_sns_ring.py +++ b/examples/Envelope/sns_ring/test_sns_ring.py @@ -5,8 +5,8 @@ import math import os import pathlib -import time import sys +import time import numpy as np import matplotlib.pyplot as plt @@ -16,12 +16,12 @@ from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.bunch_utils import collect_bunch from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.space_charge.sc2p5d import setSC2p5DAccNodes from orbit.teapot import TEAPOT_Ring from orbit.teapot import TEAPOT_MATRIX_Lattice from orbit.teapot import teapot +from orbit.teapot import BendTEAPOT from orbit.utils.consts import mass_proton sys.path.append("..") @@ -65,7 +65,7 @@ def parse_args() -> argparse.Namespace: def main(args: argparse.Namespace) -> None: path = pathlib.Path(__file__) - output_dir = os.path.join("outputs", path.stem) + output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S")) os.makedirs(output_dir, exist_ok=True) # Lattice @@ -122,7 +122,7 @@ def main(args: argparse.Namespace) -> None: if args.tilt: rot_matrix = np.identity(6) - rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=(args.tilt * math.pi)) + rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=np.radians(args.tilt)) cov_matrix = np.linalg.multi_dot([rot_matrix, cov_matrix, rot_matrix.T]) if args.mismatch_x or args.mismatch_y: @@ -162,7 +162,6 @@ def main(args: argparse.Namespace) -> None: centroid=centroid_init, intensity=args.intensity, ) - tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) history_keys = [ "rms_x", @@ -178,7 +177,7 @@ def main(args: argparse.Namespace) -> None: for turn in range(args.turns + 1): if turn > 0: - tracker.track_ring(envelope) + lattice.trackEnvelopeRing(envelope, sc=("2d" if args.sc else None)) cov_matrix = envelope.cov_matrix centroid = envelope.centroid diff --git a/examples/Envelope/sns_ring/test_sns_ring_speed.py b/examples/Envelope/sns_ring/test_sns_ring_speed.py index 8ce6a63c..124d3f75 100644 --- a/examples/Envelope/sns_ring/test_sns_ring_speed.py +++ b/examples/Envelope/sns_ring/test_sns_ring_speed.py @@ -12,11 +12,11 @@ from orbit.core.bunch import Bunch from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.space_charge.sc2p5d import setSC2p5DAccNodes from orbit.teapot import TEAPOT_Ring from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.teapot import BendTEAPOT from orbit.utils.consts import mass_proton sys.path.append("..") @@ -45,6 +45,11 @@ except: pass +for node in lattice.getNodes(): + if type(node) is BendTEAPOT: + node.setParam("ea1", 0.0) + node.setParam("ea2", 0.0) + for node in lattice.getNodes(): max_length = 1.0 if node.getLength() > max_length: @@ -84,7 +89,7 @@ cov_matrix=cov_matrix_init, intensity=args.intensity, ) -tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) +envelope_sc = "2d" if args.sc else None start_time = time.time() @@ -92,7 +97,7 @@ profiler.enable() for turn in trange(args.turns): - tracker.track_ring(envelope) + lattice.trackEnvelopeRing(envelope, sc=envelope_sc) time_per_turn = (time.time() - start_time) / args.turns diff --git a/examples/Envelope/test_env_2d_fodo.py b/examples/Envelope/test_env_2d_fodo.py index 0a4079ac..969e99ff 100644 --- a/examples/Envelope/test_env_2d_fodo.py +++ b/examples/Envelope/test_env_2d_fodo.py @@ -5,6 +5,7 @@ import math import os import pathlib +import time import numpy as np import matplotlib.pyplot as plt @@ -14,7 +15,6 @@ from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.bunch_utils import collect_bunch from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.lattice import AccLattice from orbit.lattice import AccNode from orbit.core.spacecharge import SpaceChargeCalc2p5D @@ -59,12 +59,8 @@ def parse_args() -> argparse.Namespace: def main(args: argparse.Namespace) -> None: - - # Setup - # ------------------------------------------------------------------------------ - path = pathlib.Path(__file__) - output_dir = os.path.join("outputs", path.stem) + output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S")) os.makedirs(output_dir, exist_ok=True) # Create lattice @@ -120,7 +116,7 @@ def main(args: argparse.Namespace) -> None: # Tilt if args.tilt: rot_matrix = np.identity(6) - rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=(args.tilt * math.pi)) + rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=np.radians(args.tilt)) cov_matrix = np.linalg.multi_dot([rot_matrix, cov_matrix, rot_matrix.T]) # Mismatch @@ -146,12 +142,12 @@ def main(args: argparse.Namespace) -> None: print("TRACK ENVELOPE") - tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) + envelope_sc = "2d" if args.sc else None history = {"xrms": [], "yrms": [], "xavg": [], "yavg": []} for turn in range(args.turns): if turn > 0: - tracker.track_ring(envelope) + lattice.trackEnvelopeRing(envelope, sc=envelope_sc) cov_matrix = envelope.cov_matrix centroid = envelope.centroid diff --git a/examples/Envelope/test_env_2d_fodo_speed.py b/examples/Envelope/test_env_2d_fodo_speed.py index 6882c210..e1fb5f12 100644 --- a/examples/Envelope/test_env_2d_fodo_speed.py +++ b/examples/Envelope/test_env_2d_fodo_speed.py @@ -11,7 +11,6 @@ from orbit.core.bunch import Bunch from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.space_charge.sc2p5d import setSC2p5DAccNodes from orbit.teapot import QuadTEAPOT @@ -86,7 +85,7 @@ cov_matrix=cov_matrix_init, intensity=args.intensity, ) -tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) +envelope_sc = "2d" if args.sc else None start_time = time.time() @@ -94,7 +93,7 @@ profiler.enable() for turn in trange(args.turns): - tracker.track_ring(envelope) + lattice.trackEnvelopeRing(envelope, sc=envelope_sc) time_per_turn = (time.time() - start_time) / args.turns diff --git a/examples/Envelope/test_env_3d_drift.py b/examples/Envelope/test_env_3d_drift.py index c4ec8107..992ec27d 100644 --- a/examples/Envelope/test_env_3d_drift.py +++ b/examples/Envelope/test_env_3d_drift.py @@ -9,6 +9,7 @@ import math import os import pathlib +import time import numpy as np import matplotlib.pyplot as plt @@ -19,7 +20,6 @@ from orbit.core.spacecharge import SpaceChargeCalc3D from orbit.bunch_utils import collect_bunch from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.space_charge.sc3d import setSC3DAccNodes from orbit.teapot import DriftTEAPOT from orbit.teapot import TEAPOT_Lattice @@ -42,9 +42,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--rms-y", type=float, default=0.010) parser.add_argument("--rms-z", type=float, default=0.010) - parser.add_argument("--rot-x", type=float, default=0.0) - parser.add_argument("--rot-y", type=float, default=0.0) - parser.add_argument("--rot-z", type=float, default=0.0) + parser.add_argument("--tilt-x", type=float, default=0.0) + parser.add_argument("--tilt-y", type=float, default=0.0) + parser.add_argument("--tilt-z", type=float, default=0.0) parser.add_argument("--nslice", type=int, default=10) parser.add_argument("--length", type=float, default=0.1) @@ -74,7 +74,7 @@ def build_cov_matrix_xyz( def main(args: argparse.Namespace) -> None: path = pathlib.Path(__file__) - output_dir = os.path.join("outputs", path.stem) + output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S")) os.makedirs(output_dir, exist_ok=True) # Create lattice @@ -99,7 +99,7 @@ def main(args: argparse.Namespace) -> None: cov_matrix_init = np.zeros((6, 6)) rotation_matrix = rotation_matrix_3d( - math.radians(args.rot_x), math.radians(args.rot_y), math.radians(args.rot_z) + math.radians(args.tilt_x), math.radians(args.tilt_y), math.radians(args.tilt_z) ) print(rotation_matrix) @@ -131,12 +131,12 @@ def main(args: argparse.Namespace) -> None: print("TRACK ENVELOPE") - tracker = EnvelopeTracker(lattice, sc=("3d" if args.sc else None)) + envelope_sc = "3d" if args.sc else None history = {"xrms": [], "yrms": [], "zrms": []} for turn in range(args.turns): if turn > 0: - tracker.track(envelope) + lattice.trackEnvelope(envelope, sc=envelope_sc) cov_matrix = envelope.cov_matrix diff --git a/py/orbit/envelope/__init__.py b/py/orbit/envelope/__init__.py index 7a4586ae..e2a3e379 100644 --- a/py/orbit/envelope/__init__.py +++ b/py/orbit/envelope/__init__.py @@ -1,2 +1 @@ from .envelope import Envelope -from .track import EnvelopeTracker \ No newline at end of file diff --git a/py/orbit/envelope/envelope.py b/py/orbit/envelope/envelope.py index 6c4e962c..0379818b 100644 --- a/py/orbit/envelope/envelope.py +++ b/py/orbit/envelope/envelope.py @@ -4,6 +4,7 @@ import scipy.constants import scipy.special +from orbit.core.bunch import Bunch from orbit.core.bunch import SyncParticle from orbit.utils.matrix import convert_matrix_zp_to_dE diff --git a/py/orbit/envelope/meson.build b/py/orbit/envelope/meson.build index 9be34b6b..1e9fe35c 100644 --- a/py/orbit/envelope/meson.build +++ b/py/orbit/envelope/meson.build @@ -1,7 +1,6 @@ py_sources = files([ '__init__.py', 'envelope.py', - 'track.py', 'utils.py' ]) diff --git a/py/orbit/envelope/track.py b/py/orbit/envelope/track.py deleted file mode 100644 index 0ccc4844..00000000 --- a/py/orbit/envelope/track.py +++ /dev/null @@ -1,295 +0,0 @@ -import numpy as np -import warnings - -from orbit.core.bunch import Bunch -from orbit.core.bunch import SyncParticle - -from orbit.lattice import AccNode -from orbit.lattice import AccLattice -from orbit.teapot import BendTEAPOT -from orbit.py_linac.lattice import Bend as BendLINAC - -from .envelope import Envelope - - -ENTRANCE = AccNode.ENTRANCE -BODY = AccNode.BODY -EXIT = AccNode.EXIT - -BEFORE = AccNode.BEFORE -AFTER = AccNode.AFTER - - -class EnvelopeTracker: - def __init__(self, lattice: AccLattice, sc: str | None = None) -> None: - """Constructor. - - Args: - lattice: The accelerator lattice. - sc: Envelope space charge model {"2d", "3d", None}. - """ - self.lattice = lattice - self.sc = sc - - # For pre-computing elements - self.elements = [] - self.one_turn_matrix = None - - for node in self.lattice.getNodes(): - if type(node) in (BendTEAPOT, BendLINAC): - if node.getParam("ea1") != 0.0 or node.getParam("ea2") != 0.0: - message = f"Found bend ea1 or ea2 != 0.0 ({node.getName()}.)" - message += " Nonzero edge angles are not yet supported in envelope tracking." - message += " Setting ea1 and ea2 to 0.0." - warnings.warn(message) - - node.setParam("ea1", 0.0) - node.setParam("ea2", 0.0) - - def track(self, envelope: Envelope, index_start: int = 0, index_stop: int = None) -> None: - """Track envelope through lattice. - - This is not recursive, so grandchild nodes are not tracked. - """ - nodes = self.lattice.getNodes() - nodes = nodes[index_start : index_stop] - - sync_part = envelope.sync_part - - for node_index, node in enumerate(nodes): - for child_node in node.getChildNodes(ENTRANCE): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - envelope.transform(matrix) - - for part_index in range(node.getnParts()): - for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - envelope.transform(matrix) - - matrix_sc = None - if self.sc: - length = node.getLength(part_index) - if length > 0: - if self.sc == "2d": - matrix_sc = envelope.sc_matrix_2d(length) - elif self.sc == "3d": - matrix_sc = envelope.sc_matrix_3d(length) - else: - raise ValueError - - matrix = node.getMatrix(sync_part, part_index=part_index) - if matrix is not None: - if matrix_sc is not None: - matrix = matrix @ matrix_sc - envelope.transform(matrix) - - for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - envelope.transform(matrix) - - for child_node in node.getChildNodes(EXIT): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - envelope.transform(matrix) - - def track_history(self, envelope: Envelope, index_start: int = 0, index_stop: int = None) -> dict[str, list]: - """Track and return envelope parameters vs. position in lattice.""" - history_keys = [ - "s", - "kin_energy", - "gamma", - "beta", - "mean", - "cov", - "rms_x", - "rms_y", - "rms_z", - "eps_x", - "eps_y", - ] - history = {key: [] for key in history_keys} - - def observe(envelope: Envelope) -> None: - parameters = {} - parameters["gamma"] = envelope.gamma - parameters["beta"] = envelope.beta - parameters["kin_energy"] = envelope.kin_energy - parameters["mean"] = envelope.centroid.copy() - parameters["cov"] = envelope.cov_matrix.copy() - parameters["rms_x"] = np.sqrt(parameters["cov"][0, 0]) - parameters["rms_y"] = np.sqrt(parameters["cov"][2, 2]) - parameters["rms_z"] = np.sqrt(parameters["cov"][4, 4]) - return parameters - - def update_history(envelope: Envelope, position: float) -> None: - history["s"].append(position) - parameters = observe(envelope) - for key in parameters: - history[key].append(parameters[key]) - - path_length = 0.0 - update_history(envelope, path_length) - - nodes = self.lattice.getNodes() - nodes = nodes[index_start : index_stop] - - sync_part = envelope.sync_part - - for node_index, node in enumerate(nodes): - for child_node in node.getChildNodes(ENTRANCE): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - envelope.transform(matrix) - - for part_index in range(node.getnParts()): - for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - envelope.transform(matrix) - - matrix_sc = None - if self.sc: - length = node.getLength(part_index) - if length > 0: - if self.sc == "2d": - matrix_sc = envelope.sc_matrix_2d(length) - elif self.sc == "3d": - matrix_sc = envelope.sc_matrix_3d(length) - else: - raise ValueError - - matrix = node.getMatrix(sync_part, part_index=part_index) - if matrix is not None: - if matrix_sc is not None: - matrix = matrix @ matrix_sc - envelope.transform(matrix) - - path_length += node.getLength(part_index) - update_history(envelope, path_length) - - for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - envelope.transform(matrix) - - for child_node in node.getChildNodes(EXIT): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - envelope.transform(matrix) - - return history - - def precompute_matrices(self, envelope: Envelope, index_start: int = 0, index_stop: int = None) -> None: - """Pre-compute transfer matrices for each node. - - For each node, return tuple (node, matrix). Mark space charge kicks as ("sc", length). - """ - nodes = self.lattice.getNodes() - nodes = nodes[index_start : index_stop] - - sync_part = envelope.sync_part - - self.elements = [] - for node_index, node in enumerate(nodes): - for child_node in node.getChildNodes(ENTRANCE): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - self.elements.append((child_node, matrix)) - - for part_index in range(node.getnParts()): - for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - self.elements.append((child_node, matrix)) - - if self.sc: - length = node.getLength(part_index) - if length > 0: - self.elements.append(("sc", length)) - - matrix = node.getMatrix(sync_part, part_index=part_index) - if matrix is not None: - self.elements.append((node, matrix)) - - for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - self.elements.append((node, matrix)) - - for child_node in node.getChildNodes(EXIT): - matrix = child_node.getMatrix(sync_part) - if matrix is not None: - self.elements.append((node, matrix)) - - def track_ring(self, envelope: Envelope) -> None: - """Track using pre-computed transfer matrices. - - The method assumes that all nodes are static and that there is no - change in the synchronous particle energy. In this case the matrices - can be computed once and reused on each turn. If there is no space charge, - we track using the one-turn matrix. - """ - - # Pre-compute transfer matrices on the first turn. - if not self.elements: - self.precompute_matrices(envelope) - self.one_turn_matrix = None - - # If there is no space charge, apply the one-turn transfer matrix. - if not self.sc: - if self.one_turn_matrix is None: - self.one_turn_matrix = np.identity(7) - for (node, matrix) in self.elements: - self.one_turn_matrix = matrix @ self.one_turn_matrix - return envelope.transform(self.one_turn_matrix) - - # If there is space charge, apply the matrices one-by-one. - for element in self.elements: - if element[0] == "sc": - length = element[1] - if self.sc == "2d": - envelope.transform(envelope.sc_matrix_2d(length)) - elif self.sc == "3d": - envelope.transform(envelope.sc_matrix_3d(length)) - else: - raise ValueError - else: - node, matrix = element - envelope.transform(matrix) - - def get_transfer_matrix(self, envelope: Envelope, index_start: int = 0, index_stop: int = None) -> np.ndarray: - """Return total transfer matrix (including linear space charge).""" - self.precompute_matrices(envelope, index_start, index_stop) - - if index_stop is None: - index_stop = len(self.elements) - - elements = self.elements[index_start : index_stop] - - if not self.sc: - total_matrix = np.identity(7) - for (node, matrix) in self.elements: - envelope.transform(matrix) - total_matrix = matrix @ total_matrix - return total_matrix - - total_matrix = np.identity(7) - for element in self.elements: - if element[0] == "sc": - length = element[1] - if self.sc == "2d": - matrix = envelope.sc_matrix_2d(length) - elif self.sc == "3d": - matrix = envelope.sc_matrix_3d(length) - else: - raise ValueError - envelope.transform(matrix) - total_matrix = matrix @ total_matrix - else: - node, matrix = element - envelope.transform(matrix) - total_matrix = matrix @ total_matrix - return total_matrix \ No newline at end of file diff --git a/py/orbit/lattice/AccLattice.py b/py/orbit/lattice/AccLattice.py index 84a466f8..d7ef944c 100644 --- a/py/orbit/lattice/AccLattice.py +++ b/py/orbit/lattice/AccLattice.py @@ -1,12 +1,19 @@ -import sys +from __future__ import annotations + import os +from typing import TYPE_CHECKING + +import numpy as np from ..utils import orbitFinalize from ..utils import NamedObject from ..utils import TypedObject -from ..lattice import AccActionsContainer -from ..lattice import AccNode +from .AccActionsContainer import AccActionsContainer +from .AccNode import AccNode + +if TYPE_CHECKING: + from orbit.envelope.envelope import Envelope class AccLattice(NamedObject, TypedObject): @@ -31,6 +38,9 @@ def __init__(self, name="no name"): self.__isInitialized = False self.__children = [] self.__childPositions = {} + self.__envelopeElements = [] + self.__envelopeOneTurnMatrix = None + self.__envelopeSpaceCharge = None def initialize(self): """ @@ -261,3 +271,274 @@ def trackActions(self, actionsContainer, paramsDict={}, index_start=-1, index_st paramsDict["node"] = node paramsDict["parentNode"] = self node.trackActions(actionsContainer, paramsDict) + + def _getNodesInRange(self, index_start: int = 0, index_stop: int = None): + if index_stop is None: + index_stop = len(self.__children) - 1 + return self.__children[index_start : index_stop + 1] + + def _prepareEnvelopeTracking(self) -> None: + for node in self.__children: + node_type = type(node) + is_teapot_bend = node_type.__name__ == "BendTEAPOT" and node_type.__module__ == "orbit.teapot.teapot" + is_linac_bend = node_type.__name__ == "Bend" and node_type.__module__ == "orbit.py_linac.lattice.LinacAccNodes" + if is_teapot_bend or is_linac_bend: + if node.getParam("ea1") != 0.0 or node.getParam("ea2") != 0.0: + message = f"Found bend ea1 or ea2 != 0.0 ({node.getName()}.)" + message += " Nonzero edge angles are not yet supported in envelope tracking." + message += " Please set them to zero:" + message += " `node.setParam('ea1', 0.0)`" + message += " `node.setParam('ea2', 0.0)`" + raise RuntimeError(message) + + def _getEnvelopeSpaceChargeMatrix(self, envelope: Envelope, length: float, sc: str | None) -> np.ndarray | None: + if not sc or length <= 0: + return None + if sc == "2d": + return envelope.sc_matrix_2d(length) + if sc == "3d": + return envelope.sc_matrix_3d(length) + raise ValueError(f"Invalid envelope space charge option `{sc}`") + + def setEnvelopeSpaceCharge(self, sc: str | None) -> None: + self.__envelopeSpaceCharge = sc + + def trackEnvelope( + self, + envelope: Envelope, + index_start: int = 0, + index_stop: int = None, + sc: str | None = None, + history: bool = False, + ) -> None | dict[str, list]: + """ + Track envelope through lattice. + """ + if history: + return self.trackEnvelopeHistory( + envelope, + index_start=index_start, + index_stop=index_stop, + sc=sc + ) + + self._prepareEnvelopeTracking() + self.setEnvelopeSpaceCharge(sc) + sync_part = envelope.sync_part + + for node in self._getNodesInRange(index_start, index_stop): + for child_node in node.getChildNodes(AccNode.ENTRANCE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + for part_index in range(node.getnParts()): + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.BEFORE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + matrix_sc = self._getEnvelopeSpaceChargeMatrix(envelope, node.getLength(part_index), sc) + matrix = node.getMatrix(sync_part, part_index=part_index) + if matrix is not None: + if matrix_sc is not None: + matrix = matrix @ matrix_sc + envelope.transform(matrix) + + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.AFTER): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + for child_node in node.getChildNodes(AccNode.EXIT): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + def trackEnvelopeHistory( + self, + envelope: Envelope, + index_start: int = 0, + index_stop: int = None, + sc: str | None = None, + ) -> dict[str, list]: + """ + Track envelope and return parameters vs. position in lattice. + """ + self._prepareEnvelopeTracking() + self.setEnvelopeSpaceCharge(sc) + sync_part = envelope.sync_part + + history_keys = [ + "s", + "kin_energy", + "gamma", + "beta", + "mean", + "cov", + "rms_x", + "rms_y", + "rms_z", + "eps_x", + "eps_y", + ] + history = {key: [] for key in history_keys} + + def observe(envelope: Envelope) -> dict: + parameters = {} + parameters["gamma"] = envelope.gamma + parameters["beta"] = envelope.beta + parameters["kin_energy"] = envelope.kin_energy + parameters["mean"] = envelope.centroid.copy() + parameters["cov"] = envelope.cov_matrix.copy() + parameters["rms_x"] = np.sqrt(parameters["cov"][0, 0]) + parameters["rms_y"] = np.sqrt(parameters["cov"][2, 2]) + parameters["rms_z"] = np.sqrt(parameters["cov"][4, 4]) + return parameters + + def update_history(envelope: Envelope, position: float) -> None: + history["s"].append(position) + parameters = observe(envelope) + for key in parameters: + history[key].append(parameters[key]) + + path_length = 0.0 + update_history(envelope, path_length) + + for node in self._getNodesInRange(index_start, index_stop): + for child_node in node.getChildNodes(AccNode.ENTRANCE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + for part_index in range(node.getnParts()): + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.BEFORE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + matrix_sc = self._getEnvelopeSpaceChargeMatrix(envelope, node.getLength(part_index), sc) + matrix = node.getMatrix(sync_part, part_index=part_index) + if matrix is not None: + if matrix_sc is not None: + matrix = matrix @ matrix_sc + envelope.transform(matrix) + + path_length += node.getLength(part_index) + update_history(envelope, path_length) + + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.AFTER): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + for child_node in node.getChildNodes(AccNode.EXIT): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + return history + + def precomputeEnvelopeMatrices( + self, + envelope: Envelope, + index_start: int = 0, + index_stop: int = None, + sc: str | None = None, + ) -> list: + """ + Pre-compute transfer matrices for each node. + + For each node, store tuple (node, matrix). Space charge kicks are + stored as ("sc", length). + """ + self._prepareEnvelopeTracking() + sync_part = envelope.sync_part + + self.__envelopeElements = [] + self.__envelopeOneTurnMatrix = None + self.__envelopeSpaceCharge = sc + + for node in self._getNodesInRange(index_start, index_stop): + for child_node in node.getChildNodes(AccNode.ENTRANCE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + self.__envelopeElements.append((child_node, matrix)) + + for part_index in range(node.getnParts()): + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.BEFORE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + self.__envelopeElements.append((child_node, matrix)) + + if sc: + length = node.getLength(part_index) + if length > 0: + self.__envelopeElements.append(("sc", length)) + + matrix = node.getMatrix(sync_part, part_index=part_index) + if matrix is not None: + self.__envelopeElements.append((node, matrix)) + + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.AFTER): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + self.__envelopeElements.append((node, matrix)) + + for child_node in node.getChildNodes(AccNode.EXIT): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + self.__envelopeElements.append((node, matrix)) + + return self.__envelopeElements + + def trackEnvelopeRing(self, envelope: Envelope, sc: str | None = None) -> None: + """ + Track using pre-computed transfer matrices. + + The method assumes that all nodes are static and that there is no + change in the synchronous particle energy. In this case the matrices + can be computed once and reused on each turn. If there is no space charge, + we track using the one-turn matrix. + """ + if not self.__envelopeElements or self.__envelopeSpaceCharge != sc: + self.precomputeEnvelopeMatrices(envelope, sc=sc) + + if not sc: + if self.__envelopeOneTurnMatrix is None: + self.__envelopeOneTurnMatrix = np.identity(7) + for node, matrix in self.__envelopeElements: + self.__envelopeOneTurnMatrix = matrix @ self.__envelopeOneTurnMatrix + envelope.transform(self.__envelopeOneTurnMatrix) + return + + for element in self.__envelopeElements: + if element[0] == "sc": + length = element[1] + matrix = self._getEnvelopeSpaceChargeMatrix(envelope, length, sc) + envelope.transform(matrix) + else: + node, matrix = element + envelope.transform(matrix) + + def getEnvelopeTransferMatrix( + self, + envelope: Envelope, + index_start: int = 0, + index_stop: int = None, + sc: str | None = None, + ) -> np.ndarray: + """ + Return total transfer matrix, including linear space charge when requested. + """ + elements = self.precomputeEnvelopeMatrices(envelope, index_start, index_stop, sc=sc) + + total_matrix = np.identity(7) + for element in elements: + if element[0] == "sc": + length = element[1] + matrix = self._getEnvelopeSpaceChargeMatrix(envelope, length, sc) + else: + node, matrix = element + envelope.transform(matrix) + total_matrix = matrix @ total_matrix + return total_matrix diff --git a/tests/py/orbit/test_env.py b/tests/py/orbit/test_env.py index f4367cb7..647b4855 100644 --- a/tests/py/orbit/test_env.py +++ b/tests/py/orbit/test_env.py @@ -5,7 +5,6 @@ from orbit.core.linac import MatrixRfGap from orbit.bunch_utils import collect_bunch from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.lattice import AccNode from orbit.lattice import AccLattice from orbit.py_linac.lattice import Drift @@ -82,7 +81,7 @@ def track_and_compare_rms( data[k1] = {} for k2 in ["rms", "cov"]: data[k1][k2] = {} - for k3 in ["env", "bunch"]: + for k3 in ["in", "out"]: data[k1][k2][k3] = {} # Initialize bunch @@ -106,10 +105,9 @@ def track_and_compare_rms( # Track envelope envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix) - envelope_tracker = EnvelopeTracker(lattice=lattice) data["env"]["cov"]["in"] = cov_scale * envelope.cov_matrix - envelope_tracker.track(envelope) + lattice.trackEnvelope(envelope) data["env"]["cov"]["out"] = cov_scale * envelope.cov_matrix # Compare @@ -401,10 +399,9 @@ def test_track_sublattice_no_error(): for _ in range(n): lattice.addNode(DriftTEAPOT(length=0.1)) - tracker = EnvelopeTracker(lattice) for i in range(n): - tracker.track(envelope, index_start=i) - tracker.track(envelope, index_stop=-i) + lattice.trackEnvelope(envelope, index_start=i) + lattice.trackEnvelope(envelope, index_stop=-i) def test_get_total_matrix() -> None: node = DriftTEAPOT(length=2.0, nparts=50) @@ -419,12 +416,12 @@ def test_get_total_matrix() -> None: cov_matrix = make_default_cov_matrix() envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix, intensity=1e7) - tracker = EnvelopeTracker(lattice, sc="2d") + sc = "2d" envelope_out_a = envelope.copy() - tracker.track(envelope_out_a) + lattice.trackEnvelope(envelope_out_a, sc=sc) - matrix = tracker.get_transfer_matrix(envelope.copy()) + matrix = lattice.getEnvelopeTransferMatrix(envelope.copy(), sc=sc) envelope_out_b = envelope.copy() envelope_out_b.transform(matrix) assert np.all(np.isclose(envelope_out_a.cov_matrix, envelope_out_b.cov_matrix)) From 3bdda4d5a61ac8e62ecf9b8b7b2fef3a31aba24b Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Thu, 3 Sep 2026 00:54:38 -0400 Subject: [PATCH 07/13] Add to_bunch method and ability to construct Envelope directly from Bunch --- py/orbit/envelope/envelope.py | 53 +++++++++++++++++++++++++++++++++-- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/py/orbit/envelope/envelope.py b/py/orbit/envelope/envelope.py index 0379818b..95e45d31 100644 --- a/py/orbit/envelope/envelope.py +++ b/py/orbit/envelope/envelope.py @@ -5,6 +5,7 @@ import scipy.special from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis from orbit.core.bunch import SyncParticle from orbit.utils.matrix import convert_matrix_zp_to_dE @@ -13,6 +14,27 @@ from .utils import proj_cov_matrix +def get_bunch_cov_matrix(bunch: Bunch) -> Bunch: + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(6): + cov_matrix[i, j] = cov_matrix[j, i] = twiss_calc.getCorrelation(i, j) + return cov_matrix + + +def get_bunch_centroid(bunch: Bunch) -> Bunch: + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + + centroid = np.zeros(6) + for i in range(6): + centroid[i] = twiss_calc.getAverage(i) + return centroid + + def build_diag_matrix_from_xyz_eig(eigenvectors: np.ndarray) -> np.ndarray: A = np.eye(7) for i in range(eigenvectors.shape[0]): @@ -35,11 +57,27 @@ class Envelope: def __init__( self, - sync_part: SyncParticle, + sync_part: SyncParticle = None, cov_matrix: np.ndarray = None, centroid: np.ndarray = None, intensity: float = 0.0, + bunch: Bunch = None, ) -> None: + """Constructor. + + Args: + sync_part: Synchronous particle. + cov_matrix: 6 x 6 covariance matrix. + centroid: 6 x 1 centroid vector. + intensity: Total number of particles. + bunch: If provided, the parameters above are calculated from the bunch particles. + """ + + if bunch is not None: + sync_part = bunch.getSyncParticle() + cov_matrix = get_bunch_cov_matrix(bunch) + centroid = get_bunch_centroid(bunch) + intensity = bunch.getSize() * bunch.macroSize() self.sync_part = sync_part @@ -53,7 +91,6 @@ def __init__( self.intensity = intensity self.classical_radius = get_classical_radius(self.charge, self.mass) - self.charge_sign = self.charge / abs(self.charge) # For a uniform one-dimensional distribution over length L, the standard # deviation is L * sqrt(12). This quantity is used to calculate the line @@ -111,6 +148,18 @@ def sample(self, size: int, dist: str = "kv") -> np.ndarray: particles = particles + self.centroid return particles + def to_bunch(self, size: int, dist: str = "gauss") -> Bunch: + bunch = Bunch() + bunch.mass(self.mass) + bunch.charge(self.charge) + bunch.getSyncParticle().kinEnergy(self.kin_energy) + bunch.macroSize(self.intensity / size) + + particles = self.sample(size=size, dist=dist) + for i in range(particles.shape[0]): + bunch.addParticle(*particles[i]) + return bunch + def sc_matrix_2d(self, length: float) -> np.ndarray: centroid = self.centroid cov_matrix = self.cov_matrix From ee26ce2cd760eb06a45249241ad75108d65ec7b1 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Thu, 3 Sep 2026 00:55:19 -0400 Subject: [PATCH 08/13] Fix charge sign in transfer matrix generator --- py/orbit/py_linac/lattice/LinacAccNodes.py | 1 - py/orbit/teapot/teapot.py | 3 --- 2 files changed, 4 deletions(-) diff --git a/py/orbit/py_linac/lattice/LinacAccNodes.py b/py/orbit/py_linac/lattice/LinacAccNodes.py index 4186353f..1d4e2a7e 100755 --- a/py/orbit/py_linac/lattice/LinacAccNodes.py +++ b/py/orbit/py_linac/lattice/LinacAccNodes.py @@ -742,7 +742,6 @@ def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray theta = self.getParam("theta") / (nparts - 1) if part_index == 0 or part_index == nparts - 1: theta *= 0.5 - theta *= np.sign(sync_part.charge()) return get_matrix_bend(sync_part, length=length, theta=theta) diff --git a/py/orbit/teapot/teapot.py b/py/orbit/teapot/teapot.py index 1d64ec67..cd216b8d 100644 --- a/py/orbit/teapot/teapot.py +++ b/py/orbit/teapot/teapot.py @@ -1096,7 +1096,6 @@ def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray kq = self.getParam("kq") if self.waveform: kq *= self.waveform.getStrength() - kq *= np.sign(sync_part.charge()) return get_matrix_quad(sync_part, length=length, kq=kq) @@ -1323,7 +1322,6 @@ def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray theta = self.getParam("theta") / (nparts - 1) if part_index == 0 or part_index == nparts - 1: theta *= 0.5 - theta *= np.sign(sync_part.charge()) return get_matrix_bend(sync_part, length=length, theta=theta) @@ -1710,7 +1708,6 @@ def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray return None kq = self.getParam("kq") - kq *= np.sign(sync_part.charge()) if self.waveform: kq *= self.waveform.getStrength() From 2337519a3f159de0718a76cc9adff62ac592f23c Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Thu, 3 Sep 2026 00:57:36 -0400 Subject: [PATCH 09/13] Add charge +- in tests --- tests/py/orbit/test_env.py | 259 ++++++++++++++----------------------- 1 file changed, 96 insertions(+), 163 deletions(-) diff --git a/tests/py/orbit/test_env.py b/tests/py/orbit/test_env.py index 647b4855..337c18db 100644 --- a/tests/py/orbit/test_env.py +++ b/tests/py/orbit/test_env.py @@ -1,4 +1,5 @@ import numpy as np +import pytest from orbit.core.bunch import Bunch from orbit.core.bunch import BunchTwissAnalysis @@ -63,6 +64,7 @@ def track_and_compare_rms( kin_energy: float, cov_matrix: np.ndarray, nparts: int = 100_000, + charge: float = 1.0, verbose: int = 1, ) -> dict: """Track bunch/envelope and compare rms beam sizes. @@ -87,6 +89,7 @@ def track_and_compare_rms( # Initialize bunch bunch = Bunch() bunch.mass(mass_proton) + bunch.charge(charge) sync_part = bunch.getSyncParticle() sync_part.kinEnergy(kin_energy) @@ -148,196 +151,125 @@ def make_default_cov_matrix( return np.diag(np.square([rms_x, rms_xp, rms_y, rms_yp, rms_z, rms_dE])) -def test_drift_teapot( - kin_energy: float = 0.0025, - length: float = 1.0, - cov_matrix: np.ndarray = None, - nparts: int = 6, -) -> None: - node = DriftTEAPOT(length=length, nparts=nparts) +def test_drift_teapot(): + node = DriftTEAPOT(length=1.0, nparts=6) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix) -def test_drift_linac( - kin_energy: float = 0.0025, - length: float = 1.0, - cov_matrix: np.ndarray = None, - nparts: int = 6, -) -> None: +def test_drift_linac(): node = Drift() - node.setLength(length) - node.setnParts(nparts) + node.setLength(1.0) + node.setnParts(6) nodes = [node] lattice = make_lattice(nodes) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_quad_teapot( - kin_energy: float = 0.0025, - length: float = 1.0, - kq: float = 1.0, - cov_matrix: np.ndarray = None, - nparts: int = 10, -) -> None: - node = QuadTEAPOT(length=length, kq=kq, nparts=nparts) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_quad_teapot(charge: float): + node = QuadTEAPOT(length=1.0, kq=1.0, nparts=10) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_cf_teapot( - kin_energy: float = 0.0025, - length: float = 10.0, - kq: float = 1.0, - nparts: int = 10, -) -> None: - node = ContinuousLinearFocusingTEAPOT(length=length, kq=kq, nparts=nparts) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_cf_teapot(charge: float): + node = ContinuousLinearFocusingTEAPOT(length=10.0, kq=1.0, nparts=10) lattice = make_lattice([node]) - cov_matrix = np.diag(np.square([1e-3, 0, 1e-3, 0.0, 0.0, 0.0])) - track_and_compare_rms(lattice, kin_energy, cov_matrix) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) -def test_quad_linac( - kin_energy: float = 0.0025, - length: float = 1.0, - field_grad: float = 0.23, - cov_matrix: np.ndarray = None, - nparts: int = 10, -) -> None: +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_quad_linac(charge: float): node = Quad() - node.setLength(length) - node.setnParts(nparts) - node.setParam("dB/dr", field_grad) + node.setLength(1.0) + node.setnParts(10) + node.setParam("dB/dr", 0.23) nodes = [node] lattice = make_lattice(nodes) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_bend_teapot( - kin_energy: float = 0.0025, - length: float = 1.0, - theta: float = 20.0, - cov_matrix: np.ndarray = None, - nparts: int = 5, -) -> None: - node = BendTEAPOT(length=length, theta=np.radians(theta), nparts=nparts) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_bend_teapot(charge: float): + node = BendTEAPOT(length=1.0, theta=np.radians(20.0), nparts=5) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_bend_linac( - kin_energy: float = 0.0025, - length: float = 1.0, - theta: float = 20.0, - cov_matrix: np.ndarray = None, - nparts: int = 5, -) -> None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_bend_linac(charge: float): node = Bend() - node.setLength(length) - node.setnParts(nparts) - node.setParam("theta", np.radians(theta)) + node.setLength(1.0) + node.setnParts(5) + node.setParam("theta", np.radians(20.0)) nodes = [node] lattice = make_lattice(nodes) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_kick_teapot( - kin_energy: float = 0.0025, - length: float = 0.1, - kx: float = 0.001, - ky: float = 0.001, - dE: float = 0.00001, - cov_matrix: np.ndarray = None, - nparts: int = 4, -) -> None: - node = KickTEAPOT(kx=kx, ky=ky, dE=dE, length=length, nparts=nparts) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_kick_teapot(charge: float): + node = KickTEAPOT(kx=0.001, ky=0.001, dE=0.00001, length=0.1, nparts=4) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) -def test_tilt_teapot( - kin_energy: float = 0.0025, - angle: float = 0.25 * np.pi, - cov_matrix: np.ndarray = None, -) -> None: - node = TiltTEAPOT(angle=angle) +def test_tilt_teapot(): + node = TiltTEAPOT(angle=(0.25 * np.pi)) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix) -def test_tilt_linac( - kin_energy: float = 0.0025, - angle: float = 0.25 * np.pi, - cov_matrix: np.ndarray = None, -) -> None: +def test_tilt_linac(): node = TiltElement() - node.setTiltAngle(angle) - nodes = [node] - lattice = make_lattice(nodes) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_solenoid_teapot( - kin_energy: float = 0.0025, - length: float = 2.0, - B: float = 1.0, - cov_matrix: np.ndarray = None, - nparts: int = 10, -) -> None: - node = SolenoidTEAPOT(length=length, B=B, nparts=nparts) + node.setTiltAngle(0.25 * np.pi) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_solenoid_linac( - kin_energy: float = 0.0025, - length: float = 2.0, - B: float = 1.0, - cov_matrix: np.ndarray = None, - nparts: int = 10, -) -> None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_solenoid_teapot(charge: float): + node = SolenoidTEAPOT(length=2.0, B=1.0, nparts=10) + lattice = make_lattice([node]) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_solenoid_linac(charge: float): node = Solenoid() - node.setLength(length) - node.setnParts(nparts) - node.setParam("B", B) + node.setLength(2.0) + node.setnParts(10) + node.setParam("B", 1.0) nodes = [node] lattice = make_lattice(nodes) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_rf_gap_matrix( - kin_energy: float = 0.0025, - frequency: float = 402.5e06, - E0TL: float = 0.001, - phase: float = 0.0, - charge: float = -1.0, -) -> None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_rf_gap_matrix(charge: float): + kin_energy = 0.0025 + frequency = 402.5e06 + E0TL = 0.001 + phase = 0.0 + cov_matrix = make_default_cov_matrix() bunch_in = Bunch() @@ -403,12 +335,15 @@ def test_track_sublattice_no_error(): lattice.trackEnvelope(envelope, index_start=i) lattice.trackEnvelope(envelope, index_stop=-i) -def test_get_total_matrix() -> None: + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_get_total_matrix(charge: float) -> None: node = DriftTEAPOT(length=2.0, nparts=50) lattice = make_lattice([node]) bunch = Bunch() bunch.mass(mass_proton) + bunch.charge(charge) sync_part = bunch.getSyncParticle() sync_part.kinEnergy(0.001) @@ -416,12 +351,10 @@ def test_get_total_matrix() -> None: cov_matrix = make_default_cov_matrix() envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix, intensity=1e7) - sc = "2d" - envelope_out_a = envelope.copy() - lattice.trackEnvelope(envelope_out_a, sc=sc) + lattice.trackEnvelope(envelope_out_a, sc="2d") - matrix = lattice.getEnvelopeTransferMatrix(envelope.copy(), sc=sc) + matrix = lattice.getEnvelopeTransferMatrix(envelope.copy(), sc="2d") envelope_out_b = envelope.copy() envelope_out_b.transform(matrix) assert np.all(np.isclose(envelope_out_a.cov_matrix, envelope_out_b.cov_matrix)) From 9dcb25bfcbbf0fc2ec2d77e2f14fe88d83d53b21 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Thu, 3 Sep 2026 00:58:09 -0400 Subject: [PATCH 10/13] Construct envelope directly from bunch in sns linac example --- examples/Envelope/run.sh | 4 +- examples/Envelope/sns_linac/test_sns_linac.py | 44 +++---------------- 2 files changed, 9 insertions(+), 39 deletions(-) diff --git a/examples/Envelope/run.sh b/examples/Envelope/run.sh index 31054736..078af457 100755 --- a/examples/Envelope/run.sh +++ b/examples/Envelope/run.sh @@ -15,9 +15,9 @@ python test_env_3d_drift.py --sc 1 --rms-z 0.002 --tilt-x 45.0 cd sns_linac python test_sns_linac.py --sc 0 -python test_sns_linac.py --sc 1 --dist gauss -python test_sns_linac.py --sc 1 --dist waterbag python test_sns_linac.py --sc 1 --dist kv +python test_sns_linac.py --sc 1 --dist waterbag +python test_sns_linac.py --sc 1 --dist gauss cd .. cd sns_ring diff --git a/examples/Envelope/sns_linac/test_sns_linac.py b/examples/Envelope/sns_linac/test_sns_linac.py index ee644659..11cccdd9 100755 --- a/examples/Envelope/sns_linac/test_sns_linac.py +++ b/examples/Envelope/sns_linac/test_sns_linac.py @@ -145,40 +145,6 @@ def make_bunch(args: argparse.Namespace) -> Bunch: return bunch -def get_bunch_cov_matrix(bunch: Bunch) -> Bunch: - twiss_calc = BunchTwissAnalysis() - twiss_calc.analyzeBunch(bunch) - - cov_matrix = np.zeros((6, 6)) - for i in range(6): - for j in range(6): - cov_matrix[i, j] = cov_matrix[j, i] = twiss_calc.getCorrelation(i, j) - return cov_matrix - - -def get_bunch_centroid(bunch: Bunch) -> Bunch: - twiss_calc = BunchTwissAnalysis() - twiss_calc.analyzeBunch(bunch) - - centroid = np.zeros(6) - for i in range(6): - centroid[i] = twiss_calc.getAverage(i) - return centroid - - -def make_envelope(bunch: Bunch) -> Envelope: - sync_part = bunch.getSyncParticle() - cov_matrix = get_bunch_cov_matrix(bunch) - centroid = get_bunch_centroid(bunch) - intensity = bunch.getSize() - return Envelope( - sync_part, - cov_matrix=cov_matrix, - centroid=centroid, - intensity=intensity - ) - - def main(args: argparse.Namespace) -> None: path = pathlib.Path(__file__) @@ -189,13 +155,17 @@ def main(args: argparse.Namespace) -> None: # Track envelope bunch = make_bunch(args) - envelope = make_envelope(bunch) + envelope = Envelope(bunch=bunch) lattice = make_lattice(args) lattice.trackDesignBunch(bunch) histories = {} - histories["envelope"] = lattice.trackEnvelopeHistory(envelope) + histories["envelope"] = lattice.trackEnvelope( + envelope, + history=True, + sc=("3d" if args.sc else None) + ) # Track bunch bunch = make_bunch(args) @@ -317,4 +287,4 @@ def main(args: argparse.Namespace) -> None: if __name__ == "__main__": - main(parse_args()) \ No newline at end of file + main(parse_args()) From f0ceffe62703fa1f63c929e77a225869c2c90f29 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Thu, 3 Sep 2026 01:08:26 -0400 Subject: [PATCH 11/13] Fix return types and copy cov/centroid when copying envelope --- py/orbit/envelope/envelope.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/py/orbit/envelope/envelope.py b/py/orbit/envelope/envelope.py index 95e45d31..3aeab4a9 100644 --- a/py/orbit/envelope/envelope.py +++ b/py/orbit/envelope/envelope.py @@ -1,5 +1,3 @@ -from multiprocessing.sharedctypes import Synchronized - import numpy as np import scipy.constants import scipy.special @@ -14,7 +12,7 @@ from .utils import proj_cov_matrix -def get_bunch_cov_matrix(bunch: Bunch) -> Bunch: +def get_bunch_cov_matrix(bunch: Bunch) -> np.ndarray: twiss_calc = BunchTwissAnalysis() twiss_calc.analyzeBunch(bunch) @@ -25,7 +23,7 @@ def get_bunch_cov_matrix(bunch: Bunch) -> Bunch: return cov_matrix -def get_bunch_centroid(bunch: Bunch) -> Bunch: +def get_bunch_centroid(bunch: Bunch) -> np.ndarray: twiss_calc = BunchTwissAnalysis() twiss_calc.analyzeBunch(bunch) @@ -100,8 +98,8 @@ def __init__( def copy(self): return Envelope( sync_part=self.sync_part, - cov_matrix=self.cov_matrix, - centroid=self.centroid, + cov_matrix=self.cov_matrix.copy(), + centroid=self.centroid.copy(), intensity=self.intensity ) From 171529221180bd4aafe8098cff65d64291312020 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Thu, 3 Sep 2026 01:08:50 -0400 Subject: [PATCH 12/13] Fix typo --- py/orbit/lattice/AccLattice.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/py/orbit/lattice/AccLattice.py b/py/orbit/lattice/AccLattice.py index d7ef944c..b16da010 100644 --- a/py/orbit/lattice/AccLattice.py +++ b/py/orbit/lattice/AccLattice.py @@ -482,12 +482,12 @@ def precomputeEnvelopeMatrices( for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.AFTER): matrix = child_node.getMatrix(sync_part) if matrix is not None: - self.__envelopeElements.append((node, matrix)) + self.__envelopeElements.append((child_node, matrix)) for child_node in node.getChildNodes(AccNode.EXIT): matrix = child_node.getMatrix(sync_part) if matrix is not None: - self.__envelopeElements.append((node, matrix)) + self.__envelopeElements.append((child_node, matrix)) return self.__envelopeElements From 523f3f7c6a135dfa54cf512bd0793033ada978f3 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Thu, 3 Sep 2026 01:18:33 -0400 Subject: [PATCH 13/13] Change default number of turns in sns ring example --- examples/Envelope/sns_ring/test_sns_ring_speed.py | 1 + examples/Envelope/test_env_2d_fodo_speed.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/Envelope/sns_ring/test_sns_ring_speed.py b/examples/Envelope/sns_ring/test_sns_ring_speed.py index 124d3f75..ea119c71 100644 --- a/examples/Envelope/sns_ring/test_sns_ring_speed.py +++ b/examples/Envelope/sns_ring/test_sns_ring_speed.py @@ -22,6 +22,7 @@ sys.path.append("..") from utils import gen_dist + parser = argparse.ArgumentParser() parser.add_argument("--bunch-length", type=float, default=120.0) parser.add_argument("--kin-energy", type=float, default=1.300) diff --git a/examples/Envelope/test_env_2d_fodo_speed.py b/examples/Envelope/test_env_2d_fodo_speed.py index e1fb5f12..7462d289 100644 --- a/examples/Envelope/test_env_2d_fodo_speed.py +++ b/examples/Envelope/test_env_2d_fodo_speed.py @@ -30,7 +30,7 @@ parser.add_argument("--kq", type=float, default=0.25) parser.add_argument("--nparts", type=int, default=10_000) -parser.add_argument("--turns", type=int, default=5000) +parser.add_argument("--turns", type=int, default=500) parser.add_argument("--sc", type=int, default=0) parser.add_argument("--sc-grid", type=int, default=64) args = parser.parse_args()