diff --git a/examples/Envelope/run.sh b/examples/Envelope/run.sh index 98cf808c..078af457 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 kv +python test_sns_linac.py --sc 1 --dist waterbag +python test_sns_linac.py --sc 1 --dist gauss 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 7614d54f..11cccdd9 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 - # Bunch - # -------------------------------------------------------------------------------- + rf_gaps = lattice.getRF_Gaps() + for rf_gap in rf_gaps: + rf_gap.setCppGapModel(MatrixRfGap()) + return lattice + + +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,34 @@ 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] - - 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") - - 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 main(args: argparse.Namespace) -> None: - for index, node in enumerate(lattice.getNodes()): - print(index, type(node), node.getName()) + 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) - lattice.trackDesignBunch(bunch) + random.seed(23) # Track envelope - # -------------------------------------------------------------------------------- - - 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) + bunch = make_bunch(args) + envelope = Envelope(bunch=bunch) - envelope = Envelope(bunch=bunch, cov_matrix=cov_matrix, intensity=intensity) - - 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.trackEnvelope( + envelope, + history=True, + sc=("3d" if args.sc else None) + ) # Track bunch - # -------------------------------------------------------------------------------- - + bunch = make_bunch(args) + lattice = make_lattice(args) lattice.trackDesignBunch(bunch) if args.sc: diff --git a/examples/Envelope/sns_ring/test_sns_ring.py b/examples/Envelope/sns_ring/test_sns_ring.py index 4efe4289..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: @@ -157,12 +157,11 @@ 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, ) - 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 773febc6..ea119c71 100644 --- a/examples/Envelope/sns_ring/test_sns_ring_speed.py +++ b/examples/Envelope/sns_ring/test_sns_ring_speed.py @@ -12,16 +12,17 @@ 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("..") 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) @@ -45,6 +46,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: @@ -80,11 +86,11 @@ print("ENVELOPE") envelope = Envelope( - bunch=bunch, + sync_part=sync_part, 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 +98,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 824cc3b7..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 @@ -135,7 +131,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, @@ -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 063ddfaf..7462d289 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 @@ -31,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() @@ -82,11 +81,11 @@ print("ENVELOPE") envelope = Envelope( - bunch=bunch, + sync_part=sync_part, 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 798d4448..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) @@ -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, @@ -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 003c24ab..3aeab4a9 100644 --- a/py/orbit/envelope/envelope.py +++ b/py/orbit/envelope/envelope.py @@ -5,13 +5,34 @@ 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 +def get_bunch_cov_matrix(bunch: Bunch) -> np.ndarray: + 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) -> np.ndarray: + 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]): @@ -26,7 +47,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,45 +55,40 @@ class Envelope: def __init__( self, - bunch: Bunch, + sync_part: SyncParticle = None, cov_matrix: np.ndarray = None, centroid: np.ndarray = None, intensity: float = 0.0, + bunch: Bunch = None, ) -> None: + """Constructor. - empty_bunch = Bunch() - bunch.copyEmptyBunchTo(empty_bunch) + 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. + """ - self.bunch = empty_bunch - self.sync_part = self.bunch.getSyncParticle() + 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 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) - 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 @@ -81,9 +97,9 @@ def __init__( def copy(self): return Envelope( - bunch=self.bunch, - cov_matrix=self.cov_matrix, - centroid=self.centroid, + sync_part=self.sync_part, + cov_matrix=self.cov_matrix.copy(), + centroid=self.centroid.copy(), intensity=self.intensity ) @@ -105,7 +121,7 @@ def mass(self) -> float: @property def charge(self) -> float: - return self.bunch.charge() + return self.sync_part.charge() @property def momentum(self) -> float: @@ -130,6 +146,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 diff --git a/py/orbit/envelope/matrix.py b/py/orbit/envelope/matrix.py deleted file mode 100644 index ee0cebdb..00000000 --- a/py/orbit/envelope/matrix.py +++ /dev/null @@ -1,443 +0,0 @@ -import math - -import numpy as np - -from orbit.core.bunch import Bunch -from orbit.core.bunch import SyncParticle -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.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, - MonitorTEAPOT, - FringeFieldTEAPOT, - ApertureTEAPOT, - BunchWrapTEAPOT, - TurnCounterTEAPOT, - MarkerLINAC, - FringeFieldLINAC, -] - - -def get_matrix_tilt(angle: float) -> np.ndarray: - cos_phi = math.cos(angle) - sin_phi = math.sin(angle) - - M = np.identity(7) - M[0, 0] = M[1, 1] = +cos_phi - M[0, 2] = M[1, 3] = -sin_phi - M[2, 0] = M[3, 1] = +sin_phi - M[2, 2] = M[3, 3] = +cos_phi - return M - - -def get_matrix_kick(kx: float = 0.0, ky: float = 0.0, kE: float = 0.0) -> np.ndarray: - M = np.identity(7) - M[1, -1] = kx - M[3, -1] = ky - M[5, -1] = kE - return M - - -def get_matrix_drift(envelope: Envelope, length: float) -> np.ndarray: - sync_part = envelope.sync_part - - M = np.identity(7) - M[0, 1] = length - M[2, 3] = length - M[4, 5] = length / (sync_part.gamma() ** 2) - M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) - - sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) - return M - - -def get_matrix_quad(envelope: Envelope, length: float, kq: float) -> np.ndarray: - if abs(kq) == 0: - return get_matrix_drift(envelope=envelope, length=length) - - sync_part = envelope.sync_part - - sqrt_abs_kq = math.sqrt(abs(kq)) - - M = np.identity(7) - if kq > 0: - cx = np.cos(sqrt_abs_kq * length) - sx = np.sin(sqrt_abs_kq * length) - cy = np.cosh(sqrt_abs_kq * length) - sy = np.sinh(sqrt_abs_kq * length) - M[0, 0] = cx - M[0, 1] = +sx / sqrt_abs_kq - M[1, 0] = -sx * sqrt_abs_kq - M[1, 1] = cx - M[2, 2] = cy - M[2, 3] = sy / sqrt_abs_kq - M[3, 2] = sy * sqrt_abs_kq - M[3, 3] = cy - elif kq < 0: - cx = np.cosh(sqrt_abs_kq * length) - sx = np.sinh(sqrt_abs_kq * length) - cy = np.cos(sqrt_abs_kq * length) - sy = np.sin(sqrt_abs_kq * length) - M[0, 0] = cx - M[0, 1] = sx / sqrt_abs_kq - M[1, 0] = sx * sqrt_abs_kq - M[1, 1] = cx - M[2, 2] = cy - M[2, 3] = +sy / sqrt_abs_kq - M[3, 2] = -sy * sqrt_abs_kq - M[3, 3] = cy - - M[4, 5] = length / (sync_part.gamma()**2) - M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) - - sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) - return M - - -def get_matrix_bend(envelope: Envelope, length: float, theta: float) -> np.ndarray: - sync_part = envelope.sync_part - - rho = length / theta - cx = math.cos(theta) - sx = math.sin(theta) - - M = np.identity(7) - M[0, 0] = cx - M[0, 1] = rho * sx - M[0, 5] = rho * (1.0 - cx) - M[1, 0] = -sx / rho - M[1, 1] = cx - M[1, 5] = sx - M[2, 3] = length - M[4, 0] = -sx - M[4, 1] = -rho * (1.0 - cx) - M[4, 5] = -(sync_part.beta() ** 2) * length + rho * sx - M[:5, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) - - sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) - return M - - -def get_matrix_solenoid(envelope: Envelope, length: float, B: float) -> np.ndarray: - if B == 0: - return get_matrix_drift(envelope=envelope, length=length) - - sync_part = envelope.sync_part - - phase = B * length - - V = np.identity(7) - V[:4, :4] = 0.0 - V[0, 1] = -1.0 / B - V[0, 2] = 0.5 - V[1, 0] = 0.5 * B - V[1, 3] = 1.0 - V[2, 1] = 1.0 / B - V[2, 2] = 0.5 - V[3, 0] = -0.5 * B - V[3, 3] = 1.0 - - M = np.identity(7) - M[0, 0] = +1.0 - M[1, 1] = -1.0 - M[2, 2] = math.cos(phase) - M[2, 3] = math.sin(phase) / B - M[3, 2] = math.sin(phase) * B * -1.0 - M[3, 3] = math.cos(phase) - M[4, 5] = length / (sync_part.gamma()**2) - - M = np.linalg.inv(V) @ M @ V - M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) - - sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) - return M - - -def get_matrix_cf(envelope: Envelope, length: float, kq: float) -> np.ndarray: - if kq == 0: - return get_matrix_drift(envelope=envelope, length=length) - - sync_part = envelope.sync_part - - sqrt_abs_kq = math.sqrt(abs(kq)) - - cx = math.cos(sqrt_abs_kq * length) - sx = math.sin(sqrt_abs_kq * length) - - M = np.identity(7) - M[0, 0] = M[2, 2] = cx - M[0, 1] = M[2, 3] = +sx / sqrt_abs_kq - M[1, 0] = M[3, 2] = -sx * sqrt_abs_kq - M[1, 1] = M[3, 3] = cx - M[4, 5] = length / (sync_part.gamma()**2) - M[4, 5] *= get_dp_p_coeff(sync_part) - - sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) - return M - - -def get_matrix_rf_gap(envelope: Envelope, frequency: float, E0TL: float, phase: float) -> np.ndarray: - sync_part = envelope.sync_part - - gamma = sync_part.gamma() - beta = sync_part.beta() - mass = sync_part.mass() - charge = envelope.charge - - kin_energy_in = sync_part.kinEnergy() - charge_E0TL_sin = charge * E0TL * math.sin(phase) - kin_energy_delta = charge * E0TL * math.cos(phase) - - # Calculate parameters in the center of the gap. - sync_part.momentum(sync_part.energyToMomentum(kin_energy_in + kin_energy_delta / 2.0)) - gamma_gap = sync_part.gamma() - beta_gap = sync_part.beta() - - # Move to the end of the gap. - kin_energy_out = kin_energy_in + kin_energy_delta - sync_part.momentum(sync_part.energyToMomentum(kin_energy_out)) - - # The base RF gap is simple - no phase correction. - gamma_out = sync_part.gamma() - beta_out = sync_part.beta() - prime_coeff = (beta * gamma) / (beta_out * gamma_out) - - # Wave momentum - k = 2.0 * math.pi * frequency / speed_of_light - phase_time_coeff = k / beta - - # Transverse focusing coefficient - kappa = -charge * E0TL * k / (2.0 * mass * beta_gap**2 * beta_out * gamma_gap**2 * gamma_out) - d_rp = kappa * math.sin(phase) - - M = np.eye(7) - M[5, 4] = charge_E0TL_sin * phase_time_coeff - M[4, 4] = beta_out / beta - M[1, 1] = prime_coeff - M[3, 3] = prime_coeff - M[1, 0] = d_rp - M[3, 2] = d_rp - return M - - -def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> 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. - - Args: - node: The accelerator node. - envelope: The beam envelope. - part_index: Index of the part within the node. An index of -1 returns - the transfer matrix for the entire node. - Returns: - 7 x 7 transfer matrix or None. If None, the node can be ignored during - envelope tracking. - """ - - 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(envelope=envelope, length=length) - - elif node_type is SolenoidTEAPOT: - if length <= 0: - return None - - B = node.getParam("B") - if node.waveform: - B *= node.waveform.getStrength() - B *= envelope.charge_sign - - return get_matrix_solenoid(envelope=envelope, 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) - - elif node_type is QuadTEAPOT: - if length <= 0: - return None - - kq = node.getParam("kq") - if node.waveform: - kq *= node.waveform.getStrength() - kq *= envelope.charge_sign - - return get_matrix_quad(envelope=envelope, 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 *= envelope.charge_sign - - return get_matrix_bend(envelope=envelope, 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(envelope=envelope, length=length), - ) - else: - return get_matrix_drift(envelope=envelope, length=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 *= envelope.charge_sign - if node.waveform: - kq *= node.waveform.getStrength() - - return get_matrix_cf(envelope=envelope, length=length, kq=kq) - - elif node_type is DriftLINAC: - if length <= 0: - return None - return get_matrix_drift(envelope=envelope, length=length) - - elif node_type is QuadLINAC: - if length <= 0: - return None - - brho = 3.335640952 * envelope.momentum / envelope.charge - kq = node.getParam("dB/dr") / brho - return get_matrix_quad(envelope=envelope, 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 *= envelope.charge_sign - - return get_matrix_bend(envelope=envelope, 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 - 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 * envelope.charge * length * 0.299792 / envelope.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") * envelope.charge_sign - return get_matrix_solenoid(envelope=envelope, 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() - - sync_part = envelope.sync_part - 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: - orbitFinalize("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( - envelope=envelope, - frequency=frequency, - E0TL=(E0TL * amplitude), - phase=phase, - ) - - raise NotImplementedError(str(node)) \ No newline at end of file diff --git a/py/orbit/envelope/meson.build b/py/orbit/envelope/meson.build index 90eab844..1e9fe35c 100644 --- a/py/orbit/envelope/meson.build +++ b/py/orbit/envelope/meson.build @@ -1,8 +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 deleted file mode 100644 index 16b62b9d..00000000 --- a/py/orbit/envelope/track.py +++ /dev/null @@ -1,290 +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 .matrix import get_matrix -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] - - for node_index, node in enumerate(nodes): - for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, envelope=envelope) - 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) - 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 = get_matrix(node, envelope=envelope, 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) - if matrix is not None: - envelope.transform(matrix) - - for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, envelope=envelope) - 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] - - for node_index, node in enumerate(nodes): - for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, envelope=envelope) - 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) - 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 = get_matrix(node, envelope=envelope, 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 = get_matrix(child_node, envelope=envelope) - if matrix is not None: - envelope.transform(matrix) - - for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, envelope=envelope) - 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] - - self.elements = [] - for node_index, node in enumerate(nodes): - for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, envelope=envelope) - 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) - 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 = get_matrix(node, envelope=envelope, 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) - if matrix is not None: - self.elements.append((node, matrix)) - - for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, envelope=envelope) - 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/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/lattice/AccLattice.py b/py/orbit/lattice/AccLattice.py index 84a466f8..b16da010 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((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((child_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/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..1d4e2a7e 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,18 @@ 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 + + return get_matrix_bend(sync_part, length=length, theta=theta) + class DCorrectorH(LinacMagnetNode): """ @@ -750,6 +790,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 +843,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 +948,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 +1073,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 +1103,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..cd216b8d 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,17 @@ 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() + + return get_matrix_quad(sync_part, length=length, kq=kq) + class BendTEAPOT(NodeTEAPOT): """ @@ -1248,6 +1313,18 @@ 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 + + return get_matrix_bend(sync_part, length=length, theta=theta) + class RingRFTEAPOT(NodeTEAPOT): """ @@ -1414,6 +1491,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 +1546,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 +1581,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 +1701,14 @@ 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") + 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 new file mode 100644 index 00000000..0f884f2f --- /dev/null +++ b/py/orbit/utils/matrix.py @@ -0,0 +1,247 @@ +"""Transfer matrix definitions.""" +import math + +import numpy as np + +from orbit.core.bunch import SyncParticle +from orbit.core.orbit_utils import Matrix +from orbit.utils.consts import speed_of_light + + +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) + + M = np.identity(7) + M[0, 0] = M[1, 1] = +cos_phi + M[0, 2] = M[1, 3] = -sin_phi + M[2, 0] = M[3, 1] = +sin_phi + M[2, 2] = M[3, 3] = +cos_phi + return M + + +def get_matrix_kick(kx: float = 0.0, ky: float = 0.0, kE: float = 0.0) -> np.ndarray: + M = np.identity(7) + M[1, -1] = kx + M[3, -1] = ky + M[5, -1] = kE + return M + + +def get_matrix_drift(sync_part: SyncParticle, length: float) -> np.ndarray: + M = np.identity(7) + M[0, 1] = length + M[2, 3] = length + M[4, 5] = length / (sync_part.gamma() ** 2) + M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_quad(sync_part: SyncParticle, length: float, kq: float) -> np.ndarray: + if abs(kq) == 0: + return get_matrix_drift(sync_part, length) + + sqrt_abs_kq = math.sqrt(abs(kq)) + + M = np.identity(7) + if kq > 0: + cx = np.cos(sqrt_abs_kq * length) + sx = np.sin(sqrt_abs_kq * length) + cy = np.cosh(sqrt_abs_kq * length) + sy = np.sinh(sqrt_abs_kq * length) + M[0, 0] = cx + M[0, 1] = +sx / sqrt_abs_kq + M[1, 0] = -sx * sqrt_abs_kq + M[1, 1] = cx + M[2, 2] = cy + M[2, 3] = sy / sqrt_abs_kq + M[3, 2] = sy * sqrt_abs_kq + M[3, 3] = cy + elif kq < 0: + cx = np.cosh(sqrt_abs_kq * length) + sx = np.sinh(sqrt_abs_kq * length) + cy = np.cos(sqrt_abs_kq * length) + sy = np.sin(sqrt_abs_kq * length) + M[0, 0] = cx + M[0, 1] = sx / sqrt_abs_kq + M[1, 0] = sx * sqrt_abs_kq + M[1, 1] = cx + M[2, 2] = cy + M[2, 3] = +sy / sqrt_abs_kq + M[3, 2] = -sy * sqrt_abs_kq + M[3, 3] = cy + + M[4, 5] = length / (sync_part.gamma()**2) + M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_bend(sync_part: SyncParticle, length: float, theta: float) -> np.ndarray: + rho = length / theta + cx = math.cos(theta) + sx = math.sin(theta) + + M = np.identity(7) + M[0, 0] = cx + M[0, 1] = rho * sx + M[0, 5] = rho * (1.0 - cx) + M[1, 0] = -sx / rho + M[1, 1] = cx + M[1, 5] = sx + M[2, 3] = length + M[4, 0] = -sx + M[4, 1] = -rho * (1.0 - cx) + M[4, 5] = -(sync_part.beta() ** 2) * length + rho * sx + M[:5, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_solenoid(sync_part: SyncParticle, length: float, B: float) -> np.ndarray: + if B == 0: + return get_matrix_drift(sync_part, length) + + phase = B * length + + V = np.identity(7) + V[:4, :4] = 0.0 + V[0, 1] = -1.0 / B + V[0, 2] = 0.5 + V[1, 0] = 0.5 * B + V[1, 3] = 1.0 + V[2, 1] = 1.0 / B + V[2, 2] = 0.5 + V[3, 0] = -0.5 * B + V[3, 3] = 1.0 + + M = np.identity(7) + M[0, 0] = +1.0 + M[1, 1] = -1.0 + M[2, 2] = math.cos(phase) + M[2, 3] = math.sin(phase) / B + M[3, 2] = math.sin(phase) * B * -1.0 + M[3, 3] = math.cos(phase) + M[4, 5] = length / (sync_part.gamma()**2) + + M = np.linalg.inv(V) @ M @ V + M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_cf(sync_part: SyncParticle, length: float, kq: float) -> np.ndarray: + if kq == 0: + return get_matrix_drift(sync_part, length) + + sqrt_abs_kq = math.sqrt(abs(kq)) + + cx = math.cos(sqrt_abs_kq * length) + sx = math.sin(sqrt_abs_kq * length) + + M = np.identity(7) + M[0, 0] = M[2, 2] = cx + M[0, 1] = M[2, 3] = +sx / sqrt_abs_kq + M[1, 0] = M[3, 2] = -sx * sqrt_abs_kq + M[1, 1] = M[3, 3] = cx + M[4, 5] = length / (sync_part.gamma()**2) + M[4, 5] *= get_dp_p_coeff(sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +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 = sync_part.charge() + + kin_energy_in = sync_part.kinEnergy() + charge_E0TL_sin = charge * E0TL * math.sin(phase) + kin_energy_delta = charge * E0TL * math.cos(phase) + + # Calculate parameters in the center of the gap. + sync_part.momentum(sync_part.energyToMomentum(kin_energy_in + kin_energy_delta / 2.0)) + gamma_gap = sync_part.gamma() + beta_gap = sync_part.beta() + + # Move to the end of the gap. + kin_energy_out = kin_energy_in + kin_energy_delta + sync_part.momentum(sync_part.energyToMomentum(kin_energy_out)) + + # The base RF gap is simple - no phase correction. + gamma_out = sync_part.gamma() + beta_out = sync_part.beta() + prime_coeff = (beta * gamma) / (beta_out * gamma_out) + + # Wave momentum + k = 2.0 * math.pi * frequency / speed_of_light + phase_time_coeff = k / beta + + # Transverse focusing coefficient + kappa = -charge * E0TL * k / (2.0 * mass * beta_gap**2 * beta_out * gamma_gap**2 * gamma_out) + d_rp = kappa * math.sin(phase) + + M = np.eye(7) + M[5, 4] = charge_E0TL_sin * phase_time_coeff + M[4, 4] = beta_out / beta + M[1, 1] = prime_coeff + M[3, 3] = prime_coeff + M[1, 0] = d_rp + M[3, 2] = d_rp + return M + + +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 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/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_env.py b/tests/py/orbit/test_env.py index 5275ab6c..337c18db 100644 --- a/tests/py/orbit/test_env.py +++ b/tests/py/orbit/test_env.py @@ -1,11 +1,11 @@ import numpy as np +import pytest from orbit.core.bunch import Bunch from orbit.core.bunch import BunchTwissAnalysis 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 @@ -64,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. @@ -82,13 +83,16 @@ 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 bunch = Bunch() bunch.mass(mass_proton) - bunch.getSyncParticle().kinEnergy(kin_energy) + bunch.charge(charge) + + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(kin_energy) # Track bunch particles = np.random.multivariate_normal(np.zeros(6), cov_matrix, size=nparts) @@ -103,11 +107,10 @@ 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_tracker = EnvelopeTracker(lattice=lattice) + envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix) 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 @@ -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() @@ -357,14 +289,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, @@ -386,10 +318,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() @@ -397,28 +331,30 @@ 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: +@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.getSyncParticle().kinEnergy(0.001) + bunch.charge(charge) - cov_matrix = make_default_cov_matrix() - envelope = Envelope(bunch, cov_matrix=cov_matrix, intensity=1e7) + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(0.001) - tracker = EnvelopeTracker(lattice, sc="2d") + cov_matrix = make_default_cov_matrix() + envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix, intensity=1e7) envelope_out_a = envelope.copy() - tracker.track(envelope_out_a) + lattice.trackEnvelope(envelope_out_a, sc="2d") - matrix = tracker.get_transfer_matrix(envelope.copy()) + 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)) 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()