Skip to content

xsuite_tracking

This class is used to build a Xsuite collider from a madx sequence and optics.

XsuiteTracking

XsuiteTracking class for managing particle tracking simulations.

Attributes:

Name Type Description
context_str str

The context for the simulation (e.g., "cupy", "opencl", "cpu").
device_number int

The device number for GPU contexts.
_context Context

The context object for the simulation.
beam str

The beam configuration.
distribution_file str

The file path to the particle data.
delta_max float

The maximum delta value for particles.
n_turns int

The number of turns for the simulation.
nemitt_x float

The normalized emittance in the x direction.
nemitt_y float

The normalized emittance in the y direction.

Methods:

Name Description
context

Get the context object for the simulation.
prepare_particle_distribution_for_tracking

Prepare the particle distribution for tracking.
track

Track the particles in the collider.
Source code in study_da/generate/master_classes/xsuite_tracking.py 
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
class XsuiteTracking:
    """
    XsuiteTracking class for managing particle tracking simulations.

    Attributes:
        context_str (str): The context for the simulation (e.g., "cupy", "opencl", "cpu").
        device_number (int): The device number for GPU contexts.
        _context (xo.Context): The context object for the simulation.
        beam (str): The beam configuration.
        distribution_file (str): The file path to the particle data.
        delta_max (float): The maximum delta value for particles.
        n_turns (int): The number of turns for the simulation.
        nemitt_x (float): The normalized emittance in the x direction.
        nemitt_y (float): The normalized emittance in the y direction.

    Methods:
        context: Get the context object for the simulation.
        prepare_particle_distribution_for_tracking: Prepare the particle distribution for tracking.
        track: Track the particles in the collider.
    """

    def __init__(self, configuration: dict, nemitt_x: float, nemitt_y: float) -> None:
        """
        Initialize the tracking configuration.

        Args:
            configuration (dict): A dictionary containing the configuration parameters.
                Expected keys:
                - "context": str, context string for the simulation.
                - "device_number": int, device number for the simulation.
                - "beam": str, beam type for the simulation.
                - "distribution_file": str, path to the particle file.
                - "delta_max": float, maximum delta value for the simulation.
                - "n_turns": int, number of turns for the simulation.
            nemitt_x (float): Normalized emittance in the x-plane.
            nemitt_y (float): Normalized emittance in the y-plane.
        """
        # Context parameters
        self.context_str: str = configuration["context"]
        self.device_number: int = configuration["device_number"]
        self._context = None

        # Simulation parameters
        self.beam: str = configuration["beam"]
        self.distribution_file: str = configuration["distribution_file"]
        self.path_distribution_folder_input: str = configuration["path_distribution_folder_input"]
        self.particle_path: str = f"{self.path_distribution_folder_input}/{self.distribution_file}"
        self.delta_max: float = configuration["delta_max"]
        self.n_turns: int = configuration["n_turns"]

        # Beambeam parameters
        self.nemitt_x: float = nemitt_x
        self.nemitt_y: float = nemitt_y

    @property
    def context(self) -> Any:
        """
        Returns the context for the current instance. If the context is not already set,
        it initializes the context based on the `context_str` attribute. The context can
        be one of the following:

        - "cupy": Uses `xo.ContextCupy`. If `device_number` is specified, it initializes
            the context with the given device number.
        - "opencl": Uses `xo.ContextPyopencl`.
        - "cpu": Uses `xo.ContextCpu`.
        - Any other value: Logs a warning and defaults to `xo.ContextCpu`.

        If `device_number` is specified but the context is not "cupy", a warning is logged
        indicating that the device number will be ignored.

        Returns:
            Any: The initialized context.
        """
        if self._context is None:
            if self.device_number is not None and self.context_str not in ["cupy"]:
                logging.warning("Device number will be ignored since context is not cupy")
            match self.context_str:
                case "cupy":
                    if self.device_number is not None:
                        self._context = xo.ContextCupy(device=self.device_number)
                    else:
                        self._context = xo.ContextCupy()
                case "opencl":
                    self._context = xo.ContextPyopencl()
                case "cpu":
                    self._context = xo.ContextCpu()
                case _:
                    logging.warning("Context not recognized, using cpu")
                    self._context = xo.ContextCpu()
        return self._context

    # ? I removed type hints for the output as I get an unclear linting error
    # TODO: Check the proper type hints for the output
    def prepare_particle_distribution_for_tracking(self, collider: xt.Multiline) -> tuple:
        """
        Prepare a particle distribution for tracking in the collider.

        This method reads particle data from a parquet file, processes the data to
        generate normalized amplitudes and angles, and then builds particles for
        tracking in the collider. If the context is set to use GPU, the collider
        trackers are reset and rebuilt accordingly.

        Args:
            collider (xt.Multiline): The collider object containing the beam and
                tracking information.

        Returns:
            tuple: A tuple containing:
                - xp.Particles: The particles ready for tracking.
                - np.ndarray: Array of particle IDs.
                - np.ndarray: Array of normalized amplitudes in the xy-plane.
                - np.ndarray: Array of angles in the xy-plane in radians.
        """
        # Reset the tracker to go to GPU if needed
        if self.context_str in ["cupy", "opencl"]:
            collider.discard_trackers()
            collider.build_trackers(_context=self.context)

        particle_df = pd.read_parquet(self.particle_path)

        r_vect = particle_df["normalized amplitude in xy-plane"].values
        theta_vect = particle_df["angle in xy-plane [deg]"].values * np.pi / 180  # type: ignore # [rad]

        A1_in_sigma = r_vect * np.cos(theta_vect)
        A2_in_sigma = r_vect * np.sin(theta_vect)

        particles = collider[self.beam].build_particles(
            x_norm=A1_in_sigma,
            y_norm=A2_in_sigma,
            delta=self.delta_max,
            scale_with_transverse_norm_emitt=(
                self.nemitt_x,
                self.nemitt_y,
            ),
            _context=self.context,
        )

        particle_id = particle_df.particle_id.values
        return particles, particle_id, r_vect, theta_vect

    def track(self, collider: xt.Multiline, particles: xp.Particles) -> dict:
        """
        Tracks particles through a collider for a specified number of turns and logs the elapsed time.

        Args:
            collider (xt.Multiline): The collider object containing the beamline to be tracked.
            particles (xp.Particles): The particles to be tracked.

        Returns:
            dict: A dictionary representation of the tracked particles.
        """
        # Optimize line for tracking
        collider[self.beam].optimize_for_tracking()

        # Track
        num_turns = self.n_turns
        a = time.time()
        collider[self.beam].track(particles, turn_by_turn_monitor=False, num_turns=num_turns)
        b = time.time()

        logging.info(f"Elapsed time: {b-a} s")
        logging.info(
            f"Elapsed time per particle per turn: {(b-a)/particles._capacity/num_turns*1e6} us"
        )

        return particles.to_dict()

context: Any property

Returns the context for the current instance. If the context is not already set, it initializes the context based on the context_str attribute. The context can be one of the following:

  • "cupy": Uses xo.ContextCupy. If device_number is specified, it initializes the context with the given device number.
  • "opencl": Uses xo.ContextPyopencl.
  • "cpu": Uses xo.ContextCpu.
  • Any other value: Logs a warning and defaults to xo.ContextCpu.

If device_number is specified but the context is not "cupy", a warning is logged indicating that the device number will be ignored.

Returns:

Name Type Description
Any Any

The initialized context.

__init__(configuration, nemitt_x, nemitt_y)

Initialize the tracking configuration.

Parameters:

Name Type Description Default
configuration dict

A dictionary containing the configuration parameters. Expected keys: - "context": str, context string for the simulation. - "device_number": int, device number for the simulation. - "beam": str, beam type for the simulation. - "distribution_file": str, path to the particle file. - "delta_max": float, maximum delta value for the simulation. - "n_turns": int, number of turns for the simulation.

 required
nemitt_x float

Normalized emittance in the x-plane.

 required
nemitt_y float

Normalized emittance in the y-plane.

 required
Source code in study_da/generate/master_classes/xsuite_tracking.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def __init__(self, configuration: dict, nemitt_x: float, nemitt_y: float) -> None:
    """
    Initialize the tracking configuration.

    Args:
        configuration (dict): A dictionary containing the configuration parameters.
            Expected keys:
            - "context": str, context string for the simulation.
            - "device_number": int, device number for the simulation.
            - "beam": str, beam type for the simulation.
            - "distribution_file": str, path to the particle file.
            - "delta_max": float, maximum delta value for the simulation.
            - "n_turns": int, number of turns for the simulation.
        nemitt_x (float): Normalized emittance in the x-plane.
        nemitt_y (float): Normalized emittance in the y-plane.
    """
    # Context parameters
    self.context_str: str = configuration["context"]
    self.device_number: int = configuration["device_number"]
    self._context = None

    # Simulation parameters
    self.beam: str = configuration["beam"]
    self.distribution_file: str = configuration["distribution_file"]
    self.path_distribution_folder_input: str = configuration["path_distribution_folder_input"]
    self.particle_path: str = f"{self.path_distribution_folder_input}/{self.distribution_file}"
    self.delta_max: float = configuration["delta_max"]
    self.n_turns: int = configuration["n_turns"]

    # Beambeam parameters
    self.nemitt_x: float = nemitt_x
    self.nemitt_y: float = nemitt_y

prepare_particle_distribution_for_tracking(collider)

Prepare a particle distribution for tracking in the collider.

This method reads particle data from a parquet file, processes the data to generate normalized amplitudes and angles, and then builds particles for tracking in the collider. If the context is set to use GPU, the collider trackers are reset and rebuilt accordingly.

Parameters:

Name Type Description Default
collider Multiline

The collider object containing the beam and tracking information.

 required

Returns:

Name Type Description
tuple tuple

A tuple containing: - xp.Particles: The particles ready for tracking. - np.ndarray: Array of particle IDs. - np.ndarray: Array of normalized amplitudes in the xy-plane. - np.ndarray: Array of angles in the xy-plane in radians.
Source code in study_da/generate/master_classes/xsuite_tracking.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def prepare_particle_distribution_for_tracking(self, collider: xt.Multiline) -> tuple:
    """
    Prepare a particle distribution for tracking in the collider.

    This method reads particle data from a parquet file, processes the data to
    generate normalized amplitudes and angles, and then builds particles for
    tracking in the collider. If the context is set to use GPU, the collider
    trackers are reset and rebuilt accordingly.

    Args:
        collider (xt.Multiline): The collider object containing the beam and
            tracking information.

    Returns:
        tuple: A tuple containing:
            - xp.Particles: The particles ready for tracking.
            - np.ndarray: Array of particle IDs.
            - np.ndarray: Array of normalized amplitudes in the xy-plane.
            - np.ndarray: Array of angles in the xy-plane in radians.
    """
    # Reset the tracker to go to GPU if needed
    if self.context_str in ["cupy", "opencl"]:
        collider.discard_trackers()
        collider.build_trackers(_context=self.context)

    particle_df = pd.read_parquet(self.particle_path)

    r_vect = particle_df["normalized amplitude in xy-plane"].values
    theta_vect = particle_df["angle in xy-plane [deg]"].values * np.pi / 180  # type: ignore # [rad]

    A1_in_sigma = r_vect * np.cos(theta_vect)
    A2_in_sigma = r_vect * np.sin(theta_vect)

    particles = collider[self.beam].build_particles(
        x_norm=A1_in_sigma,
        y_norm=A2_in_sigma,
        delta=self.delta_max,
        scale_with_transverse_norm_emitt=(
            self.nemitt_x,
            self.nemitt_y,
        ),
        _context=self.context,
    )

    particle_id = particle_df.particle_id.values
    return particles, particle_id, r_vect, theta_vect

track(collider, particles)

Tracks particles through a collider for a specified number of turns and logs the elapsed time.

Parameters:

Name Type Description Default
collider Multiline

The collider object containing the beamline to be tracked.

 required
particles Particles

The particles to be tracked.

 required

Returns:

Name Type Description
dict dict

A dictionary representation of the tracked particles.
Source code in study_da/generate/master_classes/xsuite_tracking.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def track(self, collider: xt.Multiline, particles: xp.Particles) -> dict:
    """
    Tracks particles through a collider for a specified number of turns and logs the elapsed time.

    Args:
        collider (xt.Multiline): The collider object containing the beamline to be tracked.
        particles (xp.Particles): The particles to be tracked.

    Returns:
        dict: A dictionary representation of the tracked particles.
    """
    # Optimize line for tracking
    collider[self.beam].optimize_for_tracking()

    # Track
    num_turns = self.n_turns
    a = time.time()
    collider[self.beam].track(particles, turn_by_turn_monitor=False, num_turns=num_turns)
    b = time.time()

    logging.info(f"Elapsed time: {b-a} s")
    logging.info(
        f"Elapsed time per particle per turn: {(b-a)/particles._capacity/num_turns*1e6} us"
    )

    return particles.to_dict()