# -*- coding: utf-8 -*-
"""
This module defines aperture elements for tracking.
"""
from abc import ABCMeta, abstractmethod
import numpy as np
from numpy.typing import NDArray
from mbtrack2.tracking.element import Element
from mbtrack2.tracking.particles import Beam, Bunch
[docs]
class Aperture(Element, metaclass=ABCMeta):
"""
Base class for aperture elements. It is not intended to be used directly.
Instead, use one of the subclasses: CircularAperture, ElipticalAperture,
RectangularAperture, or LongitudinalAperture.
"""
[docs]
def __init__(self):
super().__init__()
self.delete_particles = False
[docs]
@abstractmethod
def determine_alive(self, bunch: Bunch) -> NDArray:
"""
Determine which particles are alive based on the aperture.
This method should be overridden in subclasses.
Parameters
----------
bunch : Bunch or Beam object
The bunch of particles to check.
Returns
-------
List[bool]
A list of booleans indicating whether each particle is alive.
"""
raise NotImplementedError(
"This method should be overridden in subclasses.")
[docs]
@Element.parallel
def track(self, bunch: Bunch | Beam):
"""
Track a Beam or Bunch object through this Element.
Parameters
----------
beam : Beam or Bunch object
"""
alive = self.determine_alive(bunch)
bunch.alive[~alive] = False
if self.delete_particles:
for stat in bunch:
bunch.particles[stat] = bunch.particles[stat][alive]
bunch.mp_number = len(bunch.particles['x'])
bunch.alive = np.ones((bunch.mp_number, ), dtype=bool)
[docs]
class CircularAperture(Aperture):
"""
Circular aperture element. The particles which are outside of the circle
are 'lost' and not used in the tracking any more.
Parameters
----------
radius : float
radius of the circle in [m]
delete_particles : bool, optional
If False, they are just marked as 'not alive' and not used in the
tracking.
If True, the particles outside of the aperture are deleted.
Use with caution, if True most Monitors will not work.
Default is False.
"""
[docs]
def __init__(self, radius: float, delete_particles: bool = False):
super().__init__()
self.radius = radius
self.radius_squared = radius**2
self.delete_particles = delete_particles
[docs]
def determine_alive(self, bunch: Bunch) -> NDArray:
"""
Determine which particles are alive based on the circular aperture.
Parameters
----------
bunch : Bunch or Beam object
The bunch of particles to check.
Returns
-------
List[bool]
A list of booleans indicating whether each particle is alive.
"""
return (bunch.particles["x"]**2 + bunch.particles["y"]**2
<= self.radius_squared)
[docs]
class ElipticalAperture(Aperture):
"""
Eliptical aperture element. The particles which are outside of the elipse
are 'lost' and not used in the tracking any more.
Parameters
----------
X_radius : float
horizontal radius of the elipse in [m]
Y_radius : float
vertical radius of the elipse in [m]
delete_particles : bool, optional
If False, they are just marked as 'not alive' and not used in the
tracking.
If True, the particles outside of the aperture are deleted.
Use with caution, if True most Monitors will not work.
Default is False.
"""
[docs]
def __init__(self,
X_radius: float,
Y_radius: float,
delete_particles: bool = False):
super().__init__()
self.X_radius = X_radius
self.X_radius_squared = X_radius**2
self.Y_radius = Y_radius
self.Y_radius_squared = Y_radius**2
self.delete_particles = delete_particles
[docs]
def determine_alive(self, bunch: Bunch) -> NDArray:
"""
Determine which particles are alive based on the elliptical aperture.
Parameters
----------
bunch : Bunch or Beam object
The bunch of particles to check.
Returns
-------
List[bool]
A list of booleans indicating whether each particle is alive.
"""
return ((bunch.particles["x"]**2 / self.X_radius_squared) +
(bunch.particles["y"]**2 / self.Y_radius_squared) <= 1)
[docs]
class RectangularAperture(Aperture):
"""
Rectangular aperture element. The particles which are outside of the
rectangle are 'lost' and not used in the tracking any more.
Parameters
----------
X_right : float
right horizontal aperture of the rectangle in [m]
Y_top : float
top vertical aperture of the rectangle in [m]
X_left : float, optional
left horizontal aperture of the rectangle in [m]
Y_bottom : float, optional
bottom vertical aperture of the rectangle in [m]
delete_particles : bool, optional
If False, they are just marked as 'not alive' and not used in the
tracking.
If True, the particles outside of the aperture are deleted.
Use with caution, if True most Monitors will not work.
Default is False.
"""
[docs]
def __init__(self,
X_right: float,
Y_top: float,
X_left: float | None = None,
Y_bottom: float | None = None,
delete_particles: bool = False):
super().__init__()
self.X_right = X_right
self.X_left = X_left
self.Y_top = Y_top
self.Y_bottom = Y_bottom
self.delete_particles = delete_particles
[docs]
def determine_alive(self, bunch: Bunch) -> NDArray:
"""
Determine which particles are alive based on the rectangular aperture.
Parameters
----------
bunch : Bunch or Beam object
The bunch of particles to check.
Returns
-------
NDArray[bool]
A list of booleans indicating whether each particle is alive.
"""
if self.X_left is None:
alive_X = np.abs(bunch.particles["x"]) <= self.X_right
else:
alive_X = (bunch.particles["x"]
<= self.X_right) & (bunch.particles["x"] >= self.X_left)
if self.Y_bottom is None:
alive_Y = np.abs(bunch.particles["y"]) <= self.Y_top
else:
alive_Y = (bunch.particles["y"]
<= self.Y_top) & (bunch.particles["y"] >= self.Y_bottom)
return alive_Y & alive_X
[docs]
class LongitudinalAperture(Aperture):
"""
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].
delete_particles : bool, optional
If False, they are just marked as 'not alive' and not used in the
tracking.
If True, the particles outside of the aperture are deleted.
Use with caution, if True most Monitors will not work.
Default is False.
"""
[docs]
def __init__(self,
tau_up: float,
tau_low: float | None = None,
delete_particles: bool = False):
super().__init__()
self.tau_up: float = tau_up
self.tau_low: float = tau_low if tau_low is not None else -tau_up
self.delete_particles: bool = delete_particles
[docs]
def determine_alive(self, bunch: Bunch) -> NDArray:
"""
Determine which particles are alive based on the longitudinal aperture.
Parameters
----------
bunch : Bunch or Beam object
The bunch of particles to check.
Returns
-------
List[bool]
A list of booleans indicating whether each particle is alive.
"""
return ((bunch.particles["tau"] <= self.tau_up) &
(bunch.particles["tau"] >= self.tau_low))