diff --git a/tracking/__init__.py b/tracking/__init__.py index 05db3c85cc9817ec1f16279adae27c34bc6167ff..30f1acf38fb02d40852010e1e079302af81cc5aa 100644 --- a/tracking/__init__.py +++ b/tracking/__init__.py @@ -14,7 +14,9 @@ from mbtrack2.tracking.optics import Optics, PhyisicalModel from mbtrack2.tracking.element import (Element, LongitudinalMap, TransverseMap, SynchrotronRadiation, SkewQuadrupole) -from mbtrack2.tracking.aperture import (CircularAperture, ElipticalAperture, - RectangularAperture) +from mbtrack2.tracking.aperture import (CircularAperture, + ElipticalAperture, + RectangularAperture, + LongitudinalAperture) from mbtrack2.tracking.wakepotential import WakePotential diff --git a/tracking/aperture.py b/tracking/aperture.py index 177184563831d4ef5a93bce825127240c6d7a6c0..4ba285f0730a636cbb7de84d50624e5ae3cf87fe 100644 --- a/tracking/aperture.py +++ b/tracking/aperture.py @@ -16,13 +16,11 @@ class CircularAperture(Element): Parameters ---------- - ring : Synchrotron object radius : float radius of the circle in [m] """ - def __init__(self, ring, radius): - self.ring = ring + def __init__(self, radius): self.radius = radius self.radius_squared = radius**2 @@ -37,8 +35,8 @@ class CircularAperture(Element): ---------- bunch : Bunch or Beam object """ - alive = bunch.particles["x"]**2 + bunch.particles["y"]**2 > self.radius_squared - bunch.alive[alive] = False + alive = bunch.particles["x"]**2 + bunch.particles["y"]**2 < self.radius_squared + bunch.alive[~alive] = False class ElipticalAperture(Element): """ @@ -47,15 +45,13 @@ class ElipticalAperture(Element): Parameters ---------- - ring : Synchrotron object X_radius : float horizontal radius of the elipse in [m] Y_radius : float vertical radius of the elipse in [m] """ - def __init__(self, ring, X_radius, Y_radius): - self.ring = ring + def __init__(self, X_radius, Y_radius): self.X_radius = X_radius self.X_radius_squared = X_radius**2 self.Y_radius = Y_radius @@ -73,8 +69,8 @@ class ElipticalAperture(Element): bunch : Bunch or Beam object """ alive = (bunch.particles["x"]**2/self.X_radius_squared + - bunch.particles["y"]**2/self.Y_radius_squared > 1) - bunch.alive[alive] = False + bunch.particles["y"]**2/self.Y_radius_squared < 1) + bunch.alive[~alive] = False class RectangularAperture(Element): """ @@ -83,7 +79,6 @@ class RectangularAperture(Element): Parameters ---------- - ring : Synchrotron object X_right : float right horizontal aperture of the rectangle in [m] Y_top : float @@ -94,8 +89,7 @@ class RectangularAperture(Element): bottom vertical aperture of the rectangle in [m] """ - def __init__(self, ring, X_right, Y_top, X_left=None, Y_bottom=None): - self.ring = ring + def __init__(self, X_right, Y_top, X_left=None, Y_bottom=None): self.X_right = X_right self.X_left = X_left self.Y_top = Y_top @@ -127,4 +121,42 @@ class RectangularAperture(Element): (bunch.particles["y"] > self.Y_bottom)) alive = alive_X & alive_Y - bunch.alive[alive] = False \ No newline at end of file + bunch.alive[~alive] = False + +class LongitudinalAperture(Element): + """ + Longitudinal aperture element. The particles which are outside of the + longitudinal bounds are 'lost' and not used in the tracking any more. + + Parameters + ---------- + ring : Synchrotron object + tau_up : float + Upper longitudinal bound in [s]. + tau_low : float, optional + Lower longitudinal bound in [s]. + """ + + def __init__(self, tau_up, tau_low=None): + self.tau_up = tau_up + if tau_low is None: + self.tau_low = tau_up*-1 + else: + self.tau_low = tau_low + + @Element.parallel + def track(self, bunch): + """ + Tracking method for the element. + No bunch to bunch interaction, so written for Bunch objects and + @Element.parallel is used to handle Beam objects. + + Parameters + ---------- + bunch : Bunch or Beam object + """ + + alive = ((bunch.particles["tau"] < self.tau_up) & + (bunch.particles["tau"] > self.tau_low)) + + bunch.alive[~alive] = False \ No newline at end of file