Skip to content

generate

MadCollider

MadCollider class is responsible for setting up and managing the collider environment using MAD-X and xsuite.

Attributes:

Name Type Description
sanity_checks bool

Flag to enable or disable sanity checks.
links str

Path to the links configuration.
beam_config dict

Configuration for the beam.
optics str

Path to the optics file.
enable_imperfections bool

Flag to enable or disable imperfections.
enable_knob_synthesis bool

Flag to enable or disable knob synthesis.
rename_coupling_knobs bool

Flag to enable or disable renaming of coupling knobs.
pars_for_imperfections dict

Parameters for imperfections.
ver_lhc_run float | None

Version of LHC run.
ver_hllhc_optics float | None

Version of HL-LHC optics.
ions bool

Flag to indicate if ions are used.
phasing dict

Phasing configuration.
path_collider_file_for_configuration_as_output str

Path to save the collider.
compress bool

Flag to enable or disable compression of collider file.

Methods:

Name Description
ost

Property to get the appropriate optics specific tools.
prepare_mad_collider

Prepares the MAD-X collider environment.
build_collider

Madx, mad_b4: Madx) -> xt.Multiline: Builds the xsuite collider.
activate_RF_and_twiss

xt.Multiline) -> None: Activates RF and performs twiss analysis.
check_xsuite_lattices

xt.Line) -> None: Checks the xsuite lattices.
write_collider_to_disk

xt.Multiline) -> None: Writes the collider to disk and optionally compresses it.
clean_temporary_files

Cleans up temporary files created during the process.
Source code in study_da/generate/master_classes/mad_collider.py 
 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
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
class MadCollider:
    """
    MadCollider class is responsible for setting up and managing the collider environment using
    MAD-X and xsuite.

    Attributes:
        sanity_checks (bool): Flag to enable or disable sanity checks.
        links (str): Path to the links configuration.
        beam_config (dict): Configuration for the beam.
        optics (str): Path to the optics file.
        enable_imperfections (bool): Flag to enable or disable imperfections.
        enable_knob_synthesis (bool): Flag to enable or disable knob synthesis.
        rename_coupling_knobs (bool): Flag to enable or disable renaming of coupling knobs.
        pars_for_imperfections (dict): Parameters for imperfections.
        ver_lhc_run (float | None): Version of LHC run.
        ver_hllhc_optics (float | None): Version of HL-LHC optics.
        ions (bool): Flag to indicate if ions are used.
        phasing (dict): Phasing configuration.
        path_collider_file_for_configuration_as_output (str): Path to save the collider.
        compress (bool): Flag to enable or disable compression of collider file.

    Methods:
        ost: Property to get the appropriate optics specific tools.
        prepare_mad_collider() -> tuple[Madx, Madx]: Prepares the MAD-X collider environment.
        build_collider(mad_b1b2: Madx, mad_b4: Madx) -> xt.Multiline: Builds the xsuite collider.
        activate_RF_and_twiss(collider: xt.Multiline) -> None: Activates RF and performs twiss analysis.
        check_xsuite_lattices(line: xt.Line) -> None: Checks the xsuite lattices.
        write_collider_to_disk(collider: xt.Multiline) -> None: Writes the collider to disk and
            optionally compresses it.
        clean_temporary_files() -> None: Cleans up temporary files created during the process.
    """

    def __init__(self, configuration: dict):
        """
        Initializes the MadCollider class with the given configuration.

        Args:
            configuration (dict): A dictionary containing the following keys:
                - sanity_checks (bool): Flag to enable or disable sanity checks.
                - links (str): Path to the links configuration.
                - beam_config (dict): Configuration for the beam.
                - optics_file (str): Path to the optics file.
                - enable_imperfections (bool): Flag to enable or disable imperfections.
                - enable_knob_synthesis (bool): Flag to enable or disable knob synthesis.
                - rename_coupling_knobs (bool): Flag to enable or disable renaming of coupling
                    knobs.
                - pars_for_imperfections (dict): Parameters for imperfections.
                - ver_lhc_run (float | None): Version of the LHC run, if applicable.
                - ver_hllhc_optics (float | None): Version of the HL-LHC optics, if applicable.
                - ions (bool): Flag to indicate if ions are used.
                - phasing (dict): Configuration for phasing.
                - path_collider_file_for_configuration_as_output (str): Path to the collider.
                - compress (bool): Flag to enable or disable compression.
        """
        # Configuration variables
        self.sanity_checks: bool = configuration["sanity_checks"]
        self.links: str = configuration["links"]
        self.beam_config: dict = configuration["beam_config"]
        self.optics: str = configuration["optics_file"]
        self.enable_imperfections: bool = configuration["enable_imperfections"]
        self.enable_knob_synthesis: bool = configuration["enable_knob_synthesis"]
        self.rename_coupling_knobs: bool = configuration["rename_coupling_knobs"]
        self.pars_for_imperfections: dict = configuration["pars_for_imperfections"]
        self.ver_lhc_run: float | None = configuration["ver_lhc_run"]
        self.ver_hllhc_optics: float | None = configuration["ver_hllhc_optics"]
        self.ions: bool = configuration["ions"]
        self.phasing: dict = configuration["phasing"]

        # Optics specific tools
        self._ost = None

        # Path to disk and compression
        self.path_collider_file_for_configuration_as_output = configuration[
            "path_collider_file_for_configuration_as_output"
        ]
        self.compress = configuration["compress"]

    @property
    def ost(self) -> Any:
        """
        Determines and returns the appropriate optics-specific tools (OST) based on the
        version of HLLHC optics or LHC run configuration.

        Raises:
            ValueError: If both `ver_hllhc_optics` and `ver_lhc_run` are defined.
            ValueError: If no optics-specific tools are available for the given configuration.

        Returns:
            Any: The appropriate OST module based on the configuration.
        """
        if self._ost is None:
            # Check that version is well defined
            if self.ver_hllhc_optics is not None and self.ver_lhc_run is not None:
                raise ValueError("Only one of ver_hllhc_optics and ver_lhc_run can be defined")

            # Get the appropriate optics_specific_tools
            if self.ver_hllhc_optics is not None:
                match self.ver_hllhc_optics:
                    case 1.6:
                        self._ost = ost_hllhc16
                    case 1.3:
                        self._ost = ost_hllhc13
                    case _:
                        raise ValueError("No optics specific tools for this configuration")
            elif self.ver_lhc_run == 3.0:
                self._ost = ost_runIII_ions if self.ions else ost_runIII
            else:
                raise ValueError("No optics specific tools for the provided configuration")

        return self._ost

    def prepare_mad_collider(self) -> tuple[Madx, Madx]:
        # sourcery skip: extract-duplicate-method
        """
        Prepares the MAD-X collider environment and sequences for beam 1/2 and beam 4.

        This method performs the following steps:
        1. Creates the MAD-X environment using the provided links.
        2. Initializes MAD-X instances for beam 1/2 and beam 4 with respective command logs.
        3. Builds the sequences for both beams using the provided beam configuration.
        4. Applies the specified optics to the beam 1/2 sequence.
        5. Optionally performs sanity checks on the beam 1/2 sequence by running TWISS and checking
            the MAD-X lattices.
        6. Applies the specified optics to the beam 4 sequence.
        7. Optionally performs sanity checks on the beam 4 sequence by running TWISS and checking
            the MAD-X lattices.

        Returns:
            tuple[Madx, Madx]: A tuple containing the MAD-X instances for beam 1/2 and beam 4.
        """
        # Make mad environment
        xm.make_mad_environment(links=self.links)

        # Start mad
        mad_b1b2 = Madx(command_log="mad_collider.log")
        mad_b4 = Madx(command_log="mad_b4.log")

        # Build sequences
        self.ost.build_sequence(mad_b1b2, mylhcbeam=1, beam_config=self.beam_config)
        self.ost.build_sequence(mad_b4, mylhcbeam=4, beam_config=self.beam_config)

        # Apply optics (only for b1b2, b4 will be generated from b1b2)
        self.ost.apply_optics(mad_b1b2, optics_file=self.optics)

        if self.sanity_checks:
            mad_b1b2.use(sequence="lhcb1")
            mad_b1b2.twiss()
            self.ost.check_madx_lattices(mad_b1b2)
            mad_b1b2.use(sequence="lhcb2")
            mad_b1b2.twiss()
            self.ost.check_madx_lattices(mad_b1b2)

        # Apply optics (only for b4, just for check)
        self.ost.apply_optics(mad_b4, optics_file=self.optics)
        if self.sanity_checks:
            mad_b4.use(sequence="lhcb2")
            mad_b4.twiss()
            # ! Investigate why this is failing for run III
            try:
                self.ost.check_madx_lattices(mad_b4)
            except AssertionError:
                logging.warning("Some sanity checks have failed during the madx lattice check")

        return mad_b1b2, mad_b4

    def build_collider(self, mad_b1b2: Madx, mad_b4: Madx) -> xt.Multiline:
        """
        Build an xsuite collider using provided MAD-X sequences and configuration.

        Parameters:
        mad_b1b2 (Madx): MAD-X instance containing sequences for beam 1 and beam 2.
        mad_b4 (Madx): MAD-X instance containing sequence for beam 4.

        Returns:
        xt.Multiline: Constructed xsuite collider.

        Notes:
        - Converts `ver_lhc_run` and `ver_hllhc_optics` to float if they are not None.
        - Builds the xsuite collider with the specified sequences and configuration.
        - Optionally performs sanity checks by computing Twiss parameters for beam 1 and beam 2.
        """
        # Ensure proper types to avoid assert errors
        if self.ver_lhc_run is not None:
            self.ver_lhc_run = float(self.ver_lhc_run)
        if self.ver_hllhc_optics is not None:
            self.ver_hllhc_optics = float(self.ver_hllhc_optics)

        # Build xsuite collider
        collider = xlhc.build_xsuite_collider(
            sequence_b1=mad_b1b2.sequence.lhcb1,
            sequence_b2=mad_b1b2.sequence.lhcb2,
            sequence_b4=mad_b4.sequence.lhcb2,
            beam_config=self.beam_config,
            enable_imperfections=self.enable_imperfections,
            enable_knob_synthesis=self.enable_knob_synthesis,
            rename_coupling_knobs=self.rename_coupling_knobs,
            pars_for_imperfections=self.pars_for_imperfections,
            ver_lhc_run=self.ver_lhc_run,
            ver_hllhc_optics=self.ver_hllhc_optics,
        )
        collider.build_trackers()

        if self.sanity_checks:
            collider["lhcb1"].twiss(method="4d")
            collider["lhcb2"].twiss(method="4d")

        return collider

    def activate_RF_and_twiss(self, collider: xt.Multiline) -> None:
        """
        Activates RF and Twiss parameters for the given collider.

        This method sets the RF knobs for the collider using the values specified
        in the `phasing` attribute. It also performs sanity checks on the collider
        lattices if the `sanity_checks` attribute is set to True.

        Args:
            collider (xt.Multiline): The collider object to configure.

        Returns:
            None
        """
        # Define a RF knobs
        collider.vars["vrf400"] = self.phasing["vrf400"]
        collider.vars["lagrf400.b1"] = self.phasing["lagrf400.b1"]
        collider.vars["lagrf400.b2"] = self.phasing["lagrf400.b2"]

        if self.sanity_checks:
            for my_line in ["lhcb1", "lhcb2"]:
                self.check_xsuite_lattices(collider[my_line])

    def check_xsuite_lattices(self, line: xt.Line) -> None:
        """
        Check the Twiss parameters and tune values for a given xsuite Line object.

        This method computes the Twiss parameters for the provided `line` using the
        6-dimensional method with a specified matrix stability tolerance. It then
        prints the Twiss results at all interaction points (IPs) and the horizontal
        (Qx) and vertical (Qy) tune values.

        Args:
            line (xt.Line): The xsuite Line object for which to compute and display
                            the Twiss parameters and tune values.

        Returns:
            None
        """
        tw = line.twiss(method="6d", matrix_stability_tol=100)
        print(f"--- Now displaying Twiss result at all IPS for line {line}---")
        print(tw.rows["ip.*"])
        # print qx and qy
        print(f"--- Now displaying Qx and Qy for line {line}---")
        print(tw.qx, tw.qy)

    def write_collider_to_disk(self, collider: xt.Multiline) -> None:
        """
        Writes the collider object to disk in JSON format and optionally compresses it into a ZIP
        file.

        Args:
            collider (xt.Multiline): The collider object to be saved.

        Returns:
            None

        Raises:
            OSError: If there is an issue creating the directory or writing the file.

        Notes:
            - The method ensures that the directory specified in
                `self.path_collider_file_for_configuration_as_output` exists.
            - If `self.compress` is True, the JSON file is compressed into a ZIP file to reduce
                storage usage.
        """
        # Save collider to json, creating the folder if it does not exist
        if "/" in self.path_collider_file_for_configuration_as_output:
            os.makedirs(self.path_collider_file_for_configuration_as_output, exist_ok=True)
        collider.to_json(self.path_collider_file_for_configuration_as_output)

        # Compress the collider file to zip to ease the load on afs
        if self.compress:
            compress_and_write(self.path_collider_file_for_configuration_as_output)

    @staticmethod
    def clean_temporary_files() -> None:
        """
        Remove all the temporary files created in the process of building the collider.

        This function deletes the following files and directories:
        - "mad_collider.log"
        - "mad_b4.log"
        - "temp" directory
        - "errors"
        - "acc-models-lhc"
        """
        # Remove all the temporaty files created in the process of building collider
        os.remove("mad_collider.log")
        os.remove("mad_b4.log")
        shutil.rmtree("temp")
        os.unlink("errors")
        os.unlink("acc-models-lhc")

ost: Any property

Determines and returns the appropriate optics-specific tools (OST) based on the version of HLLHC optics or LHC run configuration.

Raises:

Type Description
ValueError

If both ver_hllhc_optics and ver_lhc_run are defined.
ValueError

If no optics-specific tools are available for the given configuration.

Returns:

Name Type Description
Any Any

The appropriate OST module based on the configuration.

__init__(configuration)

Initializes the MadCollider class with the given configuration.

Parameters:

Name Type Description Default
configuration dict

A dictionary containing the following keys: - sanity_checks (bool): Flag to enable or disable sanity checks. - links (str): Path to the links configuration. - beam_config (dict): Configuration for the beam. - optics_file (str): Path to the optics file. - enable_imperfections (bool): Flag to enable or disable imperfections. - enable_knob_synthesis (bool): Flag to enable or disable knob synthesis. - rename_coupling_knobs (bool): Flag to enable or disable renaming of coupling knobs. - pars_for_imperfections (dict): Parameters for imperfections. - ver_lhc_run (float | None): Version of the LHC run, if applicable. - ver_hllhc_optics (float | None): Version of the HL-LHC optics, if applicable. - ions (bool): Flag to indicate if ions are used. - phasing (dict): Configuration for phasing. - path_collider_file_for_configuration_as_output (str): Path to the collider. - compress (bool): Flag to enable or disable compression.

 required
Source code in study_da/generate/master_classes/mad_collider.py 
 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
def __init__(self, configuration: dict):
    """
    Initializes the MadCollider class with the given configuration.

    Args:
        configuration (dict): A dictionary containing the following keys:
            - sanity_checks (bool): Flag to enable or disable sanity checks.
            - links (str): Path to the links configuration.
            - beam_config (dict): Configuration for the beam.
            - optics_file (str): Path to the optics file.
            - enable_imperfections (bool): Flag to enable or disable imperfections.
            - enable_knob_synthesis (bool): Flag to enable or disable knob synthesis.
            - rename_coupling_knobs (bool): Flag to enable or disable renaming of coupling
                knobs.
            - pars_for_imperfections (dict): Parameters for imperfections.
            - ver_lhc_run (float | None): Version of the LHC run, if applicable.
            - ver_hllhc_optics (float | None): Version of the HL-LHC optics, if applicable.
            - ions (bool): Flag to indicate if ions are used.
            - phasing (dict): Configuration for phasing.
            - path_collider_file_for_configuration_as_output (str): Path to the collider.
            - compress (bool): Flag to enable or disable compression.
    """
    # Configuration variables
    self.sanity_checks: bool = configuration["sanity_checks"]
    self.links: str = configuration["links"]
    self.beam_config: dict = configuration["beam_config"]
    self.optics: str = configuration["optics_file"]
    self.enable_imperfections: bool = configuration["enable_imperfections"]
    self.enable_knob_synthesis: bool = configuration["enable_knob_synthesis"]
    self.rename_coupling_knobs: bool = configuration["rename_coupling_knobs"]
    self.pars_for_imperfections: dict = configuration["pars_for_imperfections"]
    self.ver_lhc_run: float | None = configuration["ver_lhc_run"]
    self.ver_hllhc_optics: float | None = configuration["ver_hllhc_optics"]
    self.ions: bool = configuration["ions"]
    self.phasing: dict = configuration["phasing"]

    # Optics specific tools
    self._ost = None

    # Path to disk and compression
    self.path_collider_file_for_configuration_as_output = configuration[
        "path_collider_file_for_configuration_as_output"
    ]
    self.compress = configuration["compress"]

activate_RF_and_twiss(collider)

Activates RF and Twiss parameters for the given collider.

This method sets the RF knobs for the collider using the values specified in the phasing attribute. It also performs sanity checks on the collider lattices if the sanity_checks attribute is set to True.

Parameters:

Name Type Description Default
collider Multiline

The collider object to configure.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/mad_collider.py 
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
def activate_RF_and_twiss(self, collider: xt.Multiline) -> None:
    """
    Activates RF and Twiss parameters for the given collider.

    This method sets the RF knobs for the collider using the values specified
    in the `phasing` attribute. It also performs sanity checks on the collider
    lattices if the `sanity_checks` attribute is set to True.

    Args:
        collider (xt.Multiline): The collider object to configure.

    Returns:
        None
    """
    # Define a RF knobs
    collider.vars["vrf400"] = self.phasing["vrf400"]
    collider.vars["lagrf400.b1"] = self.phasing["lagrf400.b1"]
    collider.vars["lagrf400.b2"] = self.phasing["lagrf400.b2"]

    if self.sanity_checks:
        for my_line in ["lhcb1", "lhcb2"]:
            self.check_xsuite_lattices(collider[my_line])

build_collider(mad_b1b2, mad_b4)

Build an xsuite collider using provided MAD-X sequences and configuration.

Parameters: mad_b1b2 (Madx): MAD-X instance containing sequences for beam 1 and beam 2. mad_b4 (Madx): MAD-X instance containing sequence for beam 4.

Returns: xt.Multiline: Constructed xsuite collider.

Notes: - Converts ver_lhc_run and ver_hllhc_optics to float if they are not None. - Builds the xsuite collider with the specified sequences and configuration. - Optionally performs sanity checks by computing Twiss parameters for beam 1 and beam 2.

Source code in study_da/generate/master_classes/mad_collider.py 
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
def build_collider(self, mad_b1b2: Madx, mad_b4: Madx) -> xt.Multiline:
    """
    Build an xsuite collider using provided MAD-X sequences and configuration.

    Parameters:
    mad_b1b2 (Madx): MAD-X instance containing sequences for beam 1 and beam 2.
    mad_b4 (Madx): MAD-X instance containing sequence for beam 4.

    Returns:
    xt.Multiline: Constructed xsuite collider.

    Notes:
    - Converts `ver_lhc_run` and `ver_hllhc_optics` to float if they are not None.
    - Builds the xsuite collider with the specified sequences and configuration.
    - Optionally performs sanity checks by computing Twiss parameters for beam 1 and beam 2.
    """
    # Ensure proper types to avoid assert errors
    if self.ver_lhc_run is not None:
        self.ver_lhc_run = float(self.ver_lhc_run)
    if self.ver_hllhc_optics is not None:
        self.ver_hllhc_optics = float(self.ver_hllhc_optics)

    # Build xsuite collider
    collider = xlhc.build_xsuite_collider(
        sequence_b1=mad_b1b2.sequence.lhcb1,
        sequence_b2=mad_b1b2.sequence.lhcb2,
        sequence_b4=mad_b4.sequence.lhcb2,
        beam_config=self.beam_config,
        enable_imperfections=self.enable_imperfections,
        enable_knob_synthesis=self.enable_knob_synthesis,
        rename_coupling_knobs=self.rename_coupling_knobs,
        pars_for_imperfections=self.pars_for_imperfections,
        ver_lhc_run=self.ver_lhc_run,
        ver_hllhc_optics=self.ver_hllhc_optics,
    )
    collider.build_trackers()

    if self.sanity_checks:
        collider["lhcb1"].twiss(method="4d")
        collider["lhcb2"].twiss(method="4d")

    return collider

check_xsuite_lattices(line)

Check the Twiss parameters and tune values for a given xsuite Line object.

This method computes the Twiss parameters for the provided line using the 6-dimensional method with a specified matrix stability tolerance. It then prints the Twiss results at all interaction points (IPs) and the horizontal (Qx) and vertical (Qy) tune values.

Parameters:

Name Type Description Default
line Line

The xsuite Line object for which to compute and display the Twiss parameters and tune values.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/mad_collider.py 
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
def check_xsuite_lattices(self, line: xt.Line) -> None:
    """
    Check the Twiss parameters and tune values for a given xsuite Line object.

    This method computes the Twiss parameters for the provided `line` using the
    6-dimensional method with a specified matrix stability tolerance. It then
    prints the Twiss results at all interaction points (IPs) and the horizontal
    (Qx) and vertical (Qy) tune values.

    Args:
        line (xt.Line): The xsuite Line object for which to compute and display
                        the Twiss parameters and tune values.

    Returns:
        None
    """
    tw = line.twiss(method="6d", matrix_stability_tol=100)
    print(f"--- Now displaying Twiss result at all IPS for line {line}---")
    print(tw.rows["ip.*"])
    # print qx and qy
    print(f"--- Now displaying Qx and Qy for line {line}---")
    print(tw.qx, tw.qy)

clean_temporary_files() staticmethod

Remove all the temporary files created in the process of building the collider.

This function deletes the following files and directories: - "mad_collider.log" - "mad_b4.log" - "temp" directory - "errors" - "acc-models-lhc"

Source code in study_da/generate/master_classes/mad_collider.py 
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
@staticmethod
def clean_temporary_files() -> None:
    """
    Remove all the temporary files created in the process of building the collider.

    This function deletes the following files and directories:
    - "mad_collider.log"
    - "mad_b4.log"
    - "temp" directory
    - "errors"
    - "acc-models-lhc"
    """
    # Remove all the temporaty files created in the process of building collider
    os.remove("mad_collider.log")
    os.remove("mad_b4.log")
    shutil.rmtree("temp")
    os.unlink("errors")
    os.unlink("acc-models-lhc")

prepare_mad_collider()

Prepares the MAD-X collider environment and sequences for beam 1/2 and beam 4.

This method performs the following steps: 1. Creates the MAD-X environment using the provided links. 2. Initializes MAD-X instances for beam 1/2 and beam 4 with respective command logs. 3. Builds the sequences for both beams using the provided beam configuration. 4. Applies the specified optics to the beam 1/2 sequence. 5. Optionally performs sanity checks on the beam 1/2 sequence by running TWISS and checking the MAD-X lattices. 6. Applies the specified optics to the beam 4 sequence. 7. Optionally performs sanity checks on the beam 4 sequence by running TWISS and checking the MAD-X lattices.

Returns:

Type Description
tuple[Madx, Madx]

tuple[Madx, Madx]: A tuple containing the MAD-X instances for beam 1/2 and beam 4.
Source code in study_da/generate/master_classes/mad_collider.py 
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
190
191
192
193
194
195
196
def prepare_mad_collider(self) -> tuple[Madx, Madx]:
    # sourcery skip: extract-duplicate-method
    """
    Prepares the MAD-X collider environment and sequences for beam 1/2 and beam 4.

    This method performs the following steps:
    1. Creates the MAD-X environment using the provided links.
    2. Initializes MAD-X instances for beam 1/2 and beam 4 with respective command logs.
    3. Builds the sequences for both beams using the provided beam configuration.
    4. Applies the specified optics to the beam 1/2 sequence.
    5. Optionally performs sanity checks on the beam 1/2 sequence by running TWISS and checking
        the MAD-X lattices.
    6. Applies the specified optics to the beam 4 sequence.
    7. Optionally performs sanity checks on the beam 4 sequence by running TWISS and checking
        the MAD-X lattices.

    Returns:
        tuple[Madx, Madx]: A tuple containing the MAD-X instances for beam 1/2 and beam 4.
    """
    # Make mad environment
    xm.make_mad_environment(links=self.links)

    # Start mad
    mad_b1b2 = Madx(command_log="mad_collider.log")
    mad_b4 = Madx(command_log="mad_b4.log")

    # Build sequences
    self.ost.build_sequence(mad_b1b2, mylhcbeam=1, beam_config=self.beam_config)
    self.ost.build_sequence(mad_b4, mylhcbeam=4, beam_config=self.beam_config)

    # Apply optics (only for b1b2, b4 will be generated from b1b2)
    self.ost.apply_optics(mad_b1b2, optics_file=self.optics)

    if self.sanity_checks:
        mad_b1b2.use(sequence="lhcb1")
        mad_b1b2.twiss()
        self.ost.check_madx_lattices(mad_b1b2)
        mad_b1b2.use(sequence="lhcb2")
        mad_b1b2.twiss()
        self.ost.check_madx_lattices(mad_b1b2)

    # Apply optics (only for b4, just for check)
    self.ost.apply_optics(mad_b4, optics_file=self.optics)
    if self.sanity_checks:
        mad_b4.use(sequence="lhcb2")
        mad_b4.twiss()
        # ! Investigate why this is failing for run III
        try:
            self.ost.check_madx_lattices(mad_b4)
        except AssertionError:
            logging.warning("Some sanity checks have failed during the madx lattice check")

    return mad_b1b2, mad_b4

write_collider_to_disk(collider)

Writes the collider object to disk in JSON format and optionally compresses it into a ZIP file.

Parameters:

Name Type Description Default
collider Multiline

The collider object to be saved.

 required

Returns:

Type Description
None

None

Raises:

Type Description
OSError

If there is an issue creating the directory or writing the file.
Notes
  • The method ensures that the directory specified in self.path_collider_file_for_configuration_as_output exists.
  • If self.compress is True, the JSON file is compressed into a ZIP file to reduce storage usage.
Source code in study_da/generate/master_classes/mad_collider.py 
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
def write_collider_to_disk(self, collider: xt.Multiline) -> None:
    """
    Writes the collider object to disk in JSON format and optionally compresses it into a ZIP
    file.

    Args:
        collider (xt.Multiline): The collider object to be saved.

    Returns:
        None

    Raises:
        OSError: If there is an issue creating the directory or writing the file.

    Notes:
        - The method ensures that the directory specified in
            `self.path_collider_file_for_configuration_as_output` exists.
        - If `self.compress` is True, the JSON file is compressed into a ZIP file to reduce
            storage usage.
    """
    # Save collider to json, creating the folder if it does not exist
    if "/" in self.path_collider_file_for_configuration_as_output:
        os.makedirs(self.path_collider_file_for_configuration_as_output, exist_ok=True)
    collider.to_json(self.path_collider_file_for_configuration_as_output)

    # Compress the collider file to zip to ease the load on afs
    if self.compress:
        compress_and_write(self.path_collider_file_for_configuration_as_output)

ParticlesDistribution

ParticlesDistribution class to generate and manage particle distributions.

Attributes:

Name Type Description
r_min int

Minimum radial distance.
r_max int

Maximum radial distance.
n_r int

Number of radial points.
n_angles int

Number of angular points.
n_split int

Number of splits for parallelization.
path_distribution_folder_output str

Path to the folder where distributions will be saved.

Methods:

Name Description
__init__

dict): Initializes the ParticlesDistribution with the given configuration.
get_radial_list

float | None = None, upper_crop: float | None = None) -> np.ndarray: Generates a list of radial distances, optionally cropped.
get_angular_list

Generates a list of angular values.
return_distribution_as_list

bool = True, lower_crop: float | None = None, upper_crop: float | None) -> list[np.ndarray]: Returns the particle distribution as a list of numpy arrays, optionally split for parallelization.
write_particle_distribution_to_disk

list[np.ndarray]) -> list[str]: Writes the particle distribution to disk in Parquet format and returns the list of file paths.
Source code in study_da/generate/master_classes/particles_distribution.py 
 23
 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
190
191
192
193
194
class ParticlesDistribution:
    """
    ParticlesDistribution class to generate and manage particle distributions.

    Attributes:
        r_min (int): Minimum radial distance.
        r_max (int): Maximum radial distance.
        n_r (int): Number of radial points.
        n_angles (int): Number of angular points.
        n_split (int): Number of splits for parallelization.
        path_distribution_folder_output (str): Path to the folder where distributions will be saved.

    Methods:
        __init__(configuration: dict):
            Initializes the ParticlesDistribution with the given configuration.

        get_radial_list(lower_crop: float | None = None, upper_crop: float | None = None)
            -> np.ndarray:
            Generates a list of radial distances, optionally cropped.

        get_angular_list() -> np.ndarray:
            Generates a list of angular values.

        return_distribution_as_list(split: bool = True, lower_crop: float | None = None,
            upper_crop: float | None) -> list[np.ndarray]:
            Returns the particle distribution as a list of numpy arrays, optionally split for
            parallelization.

        write_particle_distribution_to_disk(ll_particles: list[np.ndarray]) -> list[str]:
            Writes the particle distribution to disk in Parquet format and returns the list of file
            paths.
    """

    def __init__(self, configuration: dict):
        """
        Initialize the particle distribution with the given configuration.

        Args:
            configuration (dict): A dictionary containing the configuration parameters.
                - r_min (int): Minimum radius value.
                - r_max (int): Maximum radius value.
                - n_r (int): Number of radius points.
                - n_angles (int): Number of angle points.
                - n_split (int): Number of splits for parallelization.
                - path_distribution_folder_output (str): Path to the folder where the distribution will be
                    saved.
        """
        # Variables used to define the distribution
        self.r_min: int = configuration["r_min"]
        self.r_max: int = configuration["r_max"]
        self.n_r: int = configuration["n_r"]
        self.n_angles: int = configuration["n_angles"]

        # Variables to split the distribution for parallelization
        self.n_split: int = configuration["n_split"]

        # Variable to write the distribution to disk
        self.path_distribution_folder_output: str = configuration["path_distribution_folder_output"]

    def get_radial_list(
        self, lower_crop: float | None = None, upper_crop: float | None = None
    ) -> np.ndarray:
        """
        Generate a list of radial distances within specified bounds.

        Args:
            lower_crop (float | None): The lower bound to crop the radial distances.
                If None, no lower cropping is applied. Defaults to None.
            upper_crop (float | None): The upper bound to crop the radial distances.
                If None, no upper cropping is applied. Defaults to None.

        Returns:
            np.ndarray: An array of radial distances within the specified bounds.
        """
        radial_list = np.linspace(self.r_min, self.r_max, self.n_r, endpoint=False)
        if upper_crop:
            radial_list = radial_list[radial_list <= 7.5]
        if lower_crop:
            radial_list = radial_list[radial_list >= 2.5]
        return radial_list

    def get_angular_list(self) -> np.ndarray:
        """
        Generate a list of angular values.

        This method creates a list of angular values ranging from 0 to 90 degrees,
        excluding the first and last values. The number of angles generated is
        determined by the instance variable `self.n_angles`.

        Returns:
            numpy.ndarray: An array of angular values.
        """
        return np.linspace(0, 90, self.n_angles + 2)[1:-1]

    def return_distribution_as_list(
        self, split: bool = True, lower_crop: float | None = None, upper_crop: float | None = None
    ) -> list[np.ndarray]:
        """
        Returns the particle distribution as a list of numpy arrays.

        This method generates a particle distribution by creating a Cartesian product
        of radial and angular lists. The resulting distribution can be optionally split
        into multiple parts for parallel computation.

        Args:
            split (bool): If True, the distribution is split into multiple parts.
                Defaults to True.
            lower_crop (float | None): The lower bound for cropping the radial list.
                If None, no lower cropping is applied. Defaults to None.
            upper_crop (float | None): The upper bound for cropping the radial list.
                If None, no upper cropping is applied. Defaults to None.

        Returns:
            list[np.ndarray]: A list of numpy arrays representing the particle distribution.
                If `split` is True, the list contains multiple arrays for parallel computation.
                Otherwise, the list contains a single array.
        """
        # Get radial list and angular list
        radial_list = self.get_radial_list(lower_crop=lower_crop, upper_crop=upper_crop)
        angular_list = self.get_angular_list()

        # Define particle distribution as a cartesian product of the radial and angular lists
        l_particles = np.array(
            [
                (particle_id, ii[1], ii[0])
                for particle_id, ii in enumerate(itertools.product(angular_list, radial_list))
            ]
        )

        # Potentially split the distribution to parallelize the computation
        if split:
            return list(np.array_split(l_particles, self.n_split))

        return [l_particles]

    def write_particle_distribution_to_disk(
        self, ll_particles: list[list[np.ndarray]]
    ) -> list[str]:
        """
        Writes a list of particle distributions to disk in Parquet format.

        Args:
            ll_particles (list[list[np.ndarray]]): A list of particle distributions,
                where each distribution is a list containing particle data.

        Returns:
            list[str]: A list of file paths where the particle distributions
            have been saved.

        The method creates a directory specified by `self.path_distribution_folder_output`
        if it does not already exist. Each particle distribution is saved as a
        Parquet file in this directory. The files are named sequentially using
        a zero-padded index (e.g., '00.parquet', '01.parquet', etc.).
        """
        # Define folder to store the distributions
        os.makedirs(self.path_distribution_folder_output, exist_ok=True)

        # Write the distribution to disk
        l_path_files = []
        for idx_chunk, l_particles in enumerate(ll_particles):
            path_file = f"{self.path_distribution_folder_output}/{idx_chunk:02}.parquet"
            pd.DataFrame(
                l_particles,
                columns=[
                    "particle_id",
                    "normalized amplitude in xy-plane",
                    "angle in xy-plane [deg]",
                ],
            ).to_parquet(path_file)
            l_path_files.append(path_file)

        return l_path_files

__init__(configuration)

Initialize the particle distribution with the given configuration.

Parameters:

Name Type Description Default
configuration dict

A dictionary containing the configuration parameters. - r_min (int): Minimum radius value. - r_max (int): Maximum radius value. - n_r (int): Number of radius points. - n_angles (int): Number of angle points. - n_split (int): Number of splits for parallelization. - path_distribution_folder_output (str): Path to the folder where the distribution will be saved.

 required
Source code in study_da/generate/master_classes/particles_distribution.py 
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
def __init__(self, configuration: dict):
    """
    Initialize the particle distribution with the given configuration.

    Args:
        configuration (dict): A dictionary containing the configuration parameters.
            - r_min (int): Minimum radius value.
            - r_max (int): Maximum radius value.
            - n_r (int): Number of radius points.
            - n_angles (int): Number of angle points.
            - n_split (int): Number of splits for parallelization.
            - path_distribution_folder_output (str): Path to the folder where the distribution will be
                saved.
    """
    # Variables used to define the distribution
    self.r_min: int = configuration["r_min"]
    self.r_max: int = configuration["r_max"]
    self.n_r: int = configuration["n_r"]
    self.n_angles: int = configuration["n_angles"]

    # Variables to split the distribution for parallelization
    self.n_split: int = configuration["n_split"]

    # Variable to write the distribution to disk
    self.path_distribution_folder_output: str = configuration["path_distribution_folder_output"]

get_angular_list()

Generate a list of angular values.

This method creates a list of angular values ranging from 0 to 90 degrees, excluding the first and last values. The number of angles generated is determined by the instance variable self.n_angles.

Returns:

Type Description
ndarray

numpy.ndarray: An array of angular values.
Source code in study_da/generate/master_classes/particles_distribution.py 
104
105
106
107
108
109
110
111
112
113
114
115
def get_angular_list(self) -> np.ndarray:
    """
    Generate a list of angular values.

    This method creates a list of angular values ranging from 0 to 90 degrees,
    excluding the first and last values. The number of angles generated is
    determined by the instance variable `self.n_angles`.

    Returns:
        numpy.ndarray: An array of angular values.
    """
    return np.linspace(0, 90, self.n_angles + 2)[1:-1]

get_radial_list(lower_crop=None, upper_crop=None)

Generate a list of radial distances within specified bounds.

Parameters:

Name Type Description Default
lower_crop float | None

The lower bound to crop the radial distances. If None, no lower cropping is applied. Defaults to None.

 None
upper_crop float | None

The upper bound to crop the radial distances. If None, no upper cropping is applied. Defaults to None.

 None

Returns:

Type Description
ndarray

np.ndarray: An array of radial distances within the specified bounds.
Source code in study_da/generate/master_classes/particles_distribution.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def get_radial_list(
    self, lower_crop: float | None = None, upper_crop: float | None = None
) -> np.ndarray:
    """
    Generate a list of radial distances within specified bounds.

    Args:
        lower_crop (float | None): The lower bound to crop the radial distances.
            If None, no lower cropping is applied. Defaults to None.
        upper_crop (float | None): The upper bound to crop the radial distances.
            If None, no upper cropping is applied. Defaults to None.

    Returns:
        np.ndarray: An array of radial distances within the specified bounds.
    """
    radial_list = np.linspace(self.r_min, self.r_max, self.n_r, endpoint=False)
    if upper_crop:
        radial_list = radial_list[radial_list <= 7.5]
    if lower_crop:
        radial_list = radial_list[radial_list >= 2.5]
    return radial_list

return_distribution_as_list(split=True, lower_crop=None, upper_crop=None)

Returns the particle distribution as a list of numpy arrays.

This method generates a particle distribution by creating a Cartesian product of radial and angular lists. The resulting distribution can be optionally split into multiple parts for parallel computation.

Parameters:

Name Type Description Default
split bool

If True, the distribution is split into multiple parts. Defaults to True.

 True
lower_crop float | None

The lower bound for cropping the radial list. If None, no lower cropping is applied. Defaults to None.

 None
upper_crop float | None

The upper bound for cropping the radial list. If None, no upper cropping is applied. Defaults to None.

 None

Returns:

Type Description
list[ndarray]

list[np.ndarray]: A list of numpy arrays representing the particle distribution. If split is True, the list contains multiple arrays for parallel computation. Otherwise, the list contains a single array.
Source code in study_da/generate/master_classes/particles_distribution.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
def return_distribution_as_list(
    self, split: bool = True, lower_crop: float | None = None, upper_crop: float | None = None
) -> list[np.ndarray]:
    """
    Returns the particle distribution as a list of numpy arrays.

    This method generates a particle distribution by creating a Cartesian product
    of radial and angular lists. The resulting distribution can be optionally split
    into multiple parts for parallel computation.

    Args:
        split (bool): If True, the distribution is split into multiple parts.
            Defaults to True.
        lower_crop (float | None): The lower bound for cropping the radial list.
            If None, no lower cropping is applied. Defaults to None.
        upper_crop (float | None): The upper bound for cropping the radial list.
            If None, no upper cropping is applied. Defaults to None.

    Returns:
        list[np.ndarray]: A list of numpy arrays representing the particle distribution.
            If `split` is True, the list contains multiple arrays for parallel computation.
            Otherwise, the list contains a single array.
    """
    # Get radial list and angular list
    radial_list = self.get_radial_list(lower_crop=lower_crop, upper_crop=upper_crop)
    angular_list = self.get_angular_list()

    # Define particle distribution as a cartesian product of the radial and angular lists
    l_particles = np.array(
        [
            (particle_id, ii[1], ii[0])
            for particle_id, ii in enumerate(itertools.product(angular_list, radial_list))
        ]
    )

    # Potentially split the distribution to parallelize the computation
    if split:
        return list(np.array_split(l_particles, self.n_split))

    return [l_particles]

write_particle_distribution_to_disk(ll_particles)

Writes a list of particle distributions to disk in Parquet format.

Parameters:

Name Type Description Default
ll_particles list[list[ndarray]]

A list of particle distributions, where each distribution is a list containing particle data.

 required

Returns:

Type Description
list[str]

list[str]: A list of file paths where the particle distributions
list[str]

have been saved.

The method creates a directory specified by self.path_distribution_folder_output if it does not already exist. Each particle distribution is saved as a Parquet file in this directory. The files are named sequentially using a zero-padded index (e.g., '00.parquet', '01.parquet', etc.).

Source code in study_da/generate/master_classes/particles_distribution.py 
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
190
191
192
193
194
def write_particle_distribution_to_disk(
    self, ll_particles: list[list[np.ndarray]]
) -> list[str]:
    """
    Writes a list of particle distributions to disk in Parquet format.

    Args:
        ll_particles (list[list[np.ndarray]]): A list of particle distributions,
            where each distribution is a list containing particle data.

    Returns:
        list[str]: A list of file paths where the particle distributions
        have been saved.

    The method creates a directory specified by `self.path_distribution_folder_output`
    if it does not already exist. Each particle distribution is saved as a
    Parquet file in this directory. The files are named sequentially using
    a zero-padded index (e.g., '00.parquet', '01.parquet', etc.).
    """
    # Define folder to store the distributions
    os.makedirs(self.path_distribution_folder_output, exist_ok=True)

    # Write the distribution to disk
    l_path_files = []
    for idx_chunk, l_particles in enumerate(ll_particles):
        path_file = f"{self.path_distribution_folder_output}/{idx_chunk:02}.parquet"
        pd.DataFrame(
            l_particles,
            columns=[
                "particle_id",
                "normalized amplitude in xy-plane",
                "angle in xy-plane [deg]",
            ],
        ).to_parquet(path_file)
        l_path_files.append(path_file)

    return l_path_files

XsuiteCollider

XsuiteCollider is a class designed to handle the configuration and manipulation of a collider using the Xsuite library. It provides methods to load, configure, and tune the collider, as well as to perform luminosity leveling and beam-beam interaction setup.

Attributes:

Name Type Description
path_collider_file_for_configuration_as_input str

Path to the collider file to load.
config_beambeam dict

Configuration for beam-beam interactions.
config_knobs_and_tuning dict

Configuration for knobs and tuning.
config_lumi_leveling dict

Configuration for luminosity leveling.
config_lumi_leveling_ip1_5 dict or None

Configuration for luminosity leveling at IP1 and IP5.
config_collider dict

Configuration for the collider.
ver_hllhc_optics float

Version of the HL-LHC optics.
ver_lhc_run float

Version of the LHC run.
ions bool

Flag indicating if ions are used.
_dict_orbit_correction dict or None

Dictionary for orbit correction.
_crab bool or None

Flag indicating if crab cavities are used.
save_output_collider bool

Flag indicating if the final collider should be saved.
path_collider_file_for_tracking_as_output str

Path to save the final collider.

Methods:

Name Description
dict_orbit_correction

Property to get the dictionary for orbit correction.
load_collider

Loads the collider from a file.
install_beam_beam_wrapper

Installs beam-beam lenses in the collider.
set_knobs

Sets the knobs for the collider.
match_tune_and_chroma

Matches the tune and chromaticity of the collider.
set_filling_and_bunch_tracked

Sets the filling scheme and tracks the bunch.
compute_collision_from_scheme

Computes the number of collisions from the filling scheme.
crab

Property to get the crab cavities status.
level_all_by_separation

Levels all IPs by separation.
level_ip1_5_by_bunch_intensity

Levels IP1 and IP5 by bunch intensity.
level_ip2_8_by_separation

Levels IP2 and IP8 by separation.
add_linear_coupling

Adds linear coupling to the collider.
assert_tune_chroma_coupling

Asserts the tune, chromaticity, and coupling of the collider.
configure_beam_beam

Configures the beam-beam interactions.
record_final_luminosity

Records the final luminosity of the collider.
write_collider_to_disk

Writes the collider configuration to disk.
update_configuration_knob

Updates a specific knob in the collider.
return_fingerprint

Returns a fingerprint of the collider's configuration.
Source code in study_da/generate/master_classes/xsuite_collider.py 
 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
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
class XsuiteCollider:
    """
    XsuiteCollider is a class designed to handle the configuration and manipulation of a collider
    using the Xsuite library. It provides methods to load, configure, and tune the collider,
    as well as to perform luminosity leveling and beam-beam interaction setup.

    Attributes:
        path_collider_file_for_configuration_as_input (str): Path to the collider file to load.
        config_beambeam (dict): Configuration for beam-beam interactions.
        config_knobs_and_tuning (dict): Configuration for knobs and tuning.
        config_lumi_leveling (dict): Configuration for luminosity leveling.
        config_lumi_leveling_ip1_5 (dict or None): Configuration for luminosity leveling at IP1 and
            IP5.
        config_collider (dict): Configuration for the collider.
        ver_hllhc_optics (float): Version of the HL-LHC optics.
        ver_lhc_run (float): Version of the LHC run.
        ions (bool): Flag indicating if ions are used.
        _dict_orbit_correction (dict or None): Dictionary for orbit correction.
        _crab (bool or None): Flag indicating if crab cavities are used.
        save_output_collider (bool): Flag indicating if the final collider should be saved.
        path_collider_file_for_tracking_as_output (str): Path to save the final collider.

    Methods:
        dict_orbit_correction: Property to get the dictionary for orbit correction.
        load_collider: Loads the collider from a file.
        install_beam_beam_wrapper: Installs beam-beam lenses in the collider.
        set_knobs: Sets the knobs for the collider.
        match_tune_and_chroma: Matches the tune and chromaticity of the collider.
        set_filling_and_bunch_tracked: Sets the filling scheme and tracks the bunch.
        compute_collision_from_scheme: Computes the number of collisions from the filling scheme.
        crab: Property to get the crab cavities status.
        level_all_by_separation: Levels all IPs by separation.
        level_ip1_5_by_bunch_intensity: Levels IP1 and IP5 by bunch intensity.
        level_ip2_8_by_separation: Levels IP2 and IP8 by separation.
        add_linear_coupling: Adds linear coupling to the collider.
        assert_tune_chroma_coupling: Asserts the tune, chromaticity, and coupling of the collider.
        configure_beam_beam: Configures the beam-beam interactions.
        record_final_luminosity: Records the final luminosity of the collider.
        write_collider_to_disk: Writes the collider configuration to disk.
        update_configuration_knob: Updates a specific knob in the collider.
        return_fingerprint: Returns a fingerprint of the collider's configuration.
    """

    def __init__(
        self,
        configuration: dict,
        path_collider_file_for_configuration_as_input: str,
        ver_hllhc_optics: float,
        ver_lhc_run: float,
        ions: bool,
    ):
        """
        Initialize the XsuiteCollider class with the given configuration and parameters.

        Args:
            configuration (dict): A dictionary containing various configuration settings.
                - config_beambeam (dict): Configuration for beam-beam interactions.
                - config_knobs_and_tuning (dict): Configuration for knobs and tuning.
                - config_lumi_leveling (dict): Configuration for luminosity leveling.
                - save_output_collider (bool): Flag to save the final collider to disk.
                - path_collider_file_for_tracking_as_output (str): Path to save the final collider.
                - config_lumi_leveling_ip1_5 (optional): Configuration for luminosity leveling at
                    IP1 and IP5.
            path_collider_file_for_configuration_as_input (str): Path to the collider file.
            ver_hllhc_optics (float): Version of the HL-LHC optics.
            ver_lhc_run (float): Version of the LHC run.
            ions (bool): Flag indicating if ions are used.
        """
        # Collider file path
        self.path_collider_file_for_configuration_as_input = (
            path_collider_file_for_configuration_as_input
        )

        # Configuration variables
        self.config_beambeam: dict[str, Any] = configuration["config_beambeam"]
        self.config_knobs_and_tuning: dict[str, Any] = configuration["config_knobs_and_tuning"]
        self.config_lumi_leveling: dict[str, Any] = configuration["config_lumi_leveling"]

        # self.config_lumi_leveling_ip1_5 will be None if not present in the configuration
        self.config_lumi_leveling_ip1_5: dict[str, Any] = configuration.get(
            "config_lumi_leveling_ip1_5"
        )

        # Collider configuration
        self.config_collider: dict[str, Any] = configuration

        # Optics version (needed to select the appropriate optics specific functions)
        self.ver_hllhc_optics: float = ver_hllhc_optics
        self.ver_lhc_run: float = ver_lhc_run
        self.ions: bool = ions
        self._dict_orbit_correction: dict | None = None

        # Crab cavities
        self._crab: bool | None = None

        # Save collider to disk
        self.save_output_collider = configuration["save_output_collider"]
        self.path_collider_file_for_tracking_as_output = configuration[
            "path_collider_file_for_tracking_as_output"
        ]
        self.compress = configuration["compress"]

    @property
    def dict_orbit_correction(self) -> dict:
        """
        Generates and returns a dictionary containing orbit correction parameters.

        This method checks if the orbit correction dictionary has already been generated.
        If not, it determines the appropriate set of orbit correction parameters based on
        the version of HLLHC optics or LHC run provided.

        Returns:
            dict: A dictionary containing orbit correction parameters.

        Raises:
            ValueError: If both `ver_hllhc_optics` and `ver_lhc_run` are defined.
            ValueError: If no optics specific tools are available for the provided configuration.
        """
        if self._dict_orbit_correction is None:
            # Check that version is well defined
            if self.ver_hllhc_optics is not None and self.ver_lhc_run is not None:
                raise ValueError("Only one of ver_hllhc_optics and ver_lhc_run can be defined")

            # Get the appropriate optics_specific_tools
            if self.ver_hllhc_optics is not None:
                match self.ver_hllhc_optics:
                    case 1.6:
                        self._dict_orbit_correction = gen_corr_hllhc16()
                    case 1.3:
                        self._dict_orbit_correction = gen_corr_hllhc13()
                    case _:
                        raise ValueError("No optics specific tools for this configuration")
            elif self.ver_lhc_run == 3.0:
                self._dict_orbit_correction = (
                    gen_corr_runIII_ions() if self.ions else gen_corr_runIII()
                )
            else:
                raise ValueError("No optics specific tools for the provided configuration")

        return self._dict_orbit_correction

    @staticmethod
    def _load_collider(path_collider) -> xt.Multiline:
        """
        Load a collider configuration from a file using an external path.

        If the file path ends with ".zip", the file is uncompressed locally
        and the collider configuration is loaded from the uncompressed file.
        Otherwise, the collider configuration is loaded directly from the file.

        Returns:
            xt.Multiline: The loaded collider configuration.
        """

        # Correct collider file path if it is a zip file
        if os.path.exists(f"{path_collider}.zip") and not path_collider.endswith(".zip"):
            path_collider += ".zip"

        # Load as a json if not zip
        if not path_collider.endswith(".zip"):
            return xt.Multiline.from_json(path_collider)

        # Uncompress file locally
        logging.info(f"Unzipping {path_collider}")
        with ZipFile(path_collider, "r") as zip_ref:
            zip_ref.extractall()
        final_path = os.path.basename(path_collider).replace(".zip", "")
        return xt.Multiline.from_json(final_path)

    def load_collider(self) -> xt.Multiline:
        """
        Load a collider configuration from a file.

        If the file path ends with ".zip", the file is uncompressed locally
        and the collider configuration is loaded from the uncompressed file.
        Otherwise, the collider configuration is loaded directly from the file.

        Returns:
            xt.Multiline: The loaded collider configuration.
        """
        return self._load_collider(self.path_collider_file_for_configuration_as_input)

    def install_beam_beam_wrapper(self, collider: xt.Multiline) -> None:
        """
        This method installs beam-beam interactions in the collider with the specified
        parameters. The beam-beam lenses are initially inactive and not configured.

        Args:
            collider (xt.Multiline): The collider object where the beam-beam interactions
                will be installed.

        Returns:
            None
        """
        # Install beam-beam lenses (inactive and not configured)
        collider.install_beambeam_interactions(
            clockwise_line="lhcb1",
            anticlockwise_line="lhcb2",
            ip_names=["ip1", "ip2", "ip5", "ip8"],
            delay_at_ips_slots=[0, 891, 0, 2670],
            num_long_range_encounters_per_side=self.config_beambeam[
                "num_long_range_encounters_per_side"
            ],
            num_slices_head_on=self.config_beambeam["num_slices_head_on"],
            harmonic_number=35640,
            bunch_spacing_buckets=self.config_beambeam["bunch_spacing_buckets"],
            sigmaz=self.config_beambeam["sigma_z"],
        )

    def set_knobs(self, collider: xt.Multiline) -> None:
        """
        Set all knobs for the collider, including crossing angles, dispersion correction,
        RF, crab cavities, experimental magnets, etc.

        Args:
            collider (xt.Multiline): The collider object to which the knob settings will be applied.

        Returns:
            None
        """
        # Set all knobs (crossing angles, dispersion correction, rf, crab cavities,
        # experimental magnets, etc.)
        for kk, vv in self.config_knobs_and_tuning["knob_settings"].items():
            collider.vars[kk] = vv

        # Crab fix (if needed)
        if self.ver_hllhc_optics is not None and self.ver_hllhc_optics == 1.3:
            apply_crab_fix(collider, self.config_knobs_and_tuning)

    def match_tune_and_chroma(
        self, collider: xt.Multiline, match_linear_coupling_to_zero: bool = True
    ) -> None:
        """
        This method adjusts the tune and chromaticity of the specified collider lines
        ("lhcb1" and "lhcb2") to the target values defined in the configuration. It also
        optionally matches the linear coupling to zero.

        Args:
            collider (xt.Multiline): The collider object containing the lines to be tuned.
            match_linear_coupling_to_zero (bool, optional): If True, linear coupling will be
                matched to zero. Defaults to True.

        Returns:
            None
        """
        for line_name in ["lhcb1", "lhcb2"]:
            knob_names = self.config_knobs_and_tuning["knob_names"][line_name]

            targets = {
                "qx": self.config_knobs_and_tuning["qx"][line_name],
                "qy": self.config_knobs_and_tuning["qy"][line_name],
                "dqx": self.config_knobs_and_tuning["dqx"][line_name],
                "dqy": self.config_knobs_and_tuning["dqy"][line_name],
            }

            xm.machine_tuning(
                line=collider[line_name],
                enable_closed_orbit_correction=True,
                enable_linear_coupling_correction=match_linear_coupling_to_zero,
                enable_tune_correction=True,
                enable_chromaticity_correction=True,
                knob_names=knob_names,
                targets=targets,
                line_co_ref=collider[f"{line_name}_co_ref"],
                co_corr_config=self.dict_orbit_correction[line_name],
            )

    def set_filling_and_bunch_tracked(self, ask_worst_bunch: bool = False) -> None:
        """
        Sets the filling scheme and determines the bunch to be tracked for beam-beam interactions.

        This method performs the following steps:
        1. Retrieves the filling scheme path from the configuration.
        2. Checks if the filling scheme path needs to be obtained from the template schemes.
        3. Loads and verifies the filling scheme, potentially converting it if necessary.
        4. Updates the configuration with the correct filling scheme path.
        5. Determines the number of long-range encounters to consider.
        6. If the bunch number for beam 1 is not provided, it identifies the bunch with the largest
        number of long-range interactions.
           - If `ask_worst_bunch` is True, prompts the user to confirm or provide a bunch number.
           - Otherwise, automatically selects the worst bunch.
        7. If the bunch number for beam 2 is not provided, it automatically selects the worst bunch.

        Args:
            ask_worst_bunch (bool): If True, prompts the user to confirm or provide the bunch number
                for beam 1. Defaults to False.

        Returns:
            None
        """
        # Get the filling scheme path
        filling_scheme_path = self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"]

        # Check if the filling scheme path must be obtained from the template schemes
        scheme_folder = (
            pathlib.Path(__file__).parent.parent.parent.resolve().joinpath("assets/filling_schemes")
        )
        if filling_scheme_path in os.listdir(scheme_folder):
            filling_scheme_path = str(scheme_folder.joinpath(filling_scheme_path))
            self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"] = filling_scheme_path

        # Load and check filling scheme, potentially convert it
        filling_scheme_path = load_and_check_filling_scheme(filling_scheme_path)

        # Correct filling scheme in config, as it might have been converted
        self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"] = filling_scheme_path

        # Get number of LR to consider
        n_LR = self.config_beambeam["num_long_range_encounters_per_side"]["ip1"]

        # If the bunch number is None, the bunch with the largest number of long-range interactions is used
        if self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"] is None:
            # Case the bunch number has not been provided
            worst_bunch_b1 = get_worst_bunch(
                filling_scheme_path, number_of_LR_to_consider=n_LR, beam="beam_1"
            )
            if ask_worst_bunch:
                while self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"] is None:
                    bool_inp = input(
                        "The bunch number for beam 1 has not been provided. Do you want to use the"
                        " bunch with the largest number of long-range interactions? It is the bunch"
                        " number " + str(worst_bunch_b1) + " (y/n): "
                    )
                    if bool_inp == "y":
                        self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"] = (
                            worst_bunch_b1
                        )
                    elif bool_inp == "n":
                        self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"] = int(
                            input("Please enter the bunch number for beam 1: ")
                        )
            else:
                self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"] = worst_bunch_b1

        if self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b2"] is None:
            worst_bunch_b2 = get_worst_bunch(
                filling_scheme_path, number_of_LR_to_consider=n_LR, beam="beam_2"
            )
            # For beam 2, just select the worst bunch by default
            self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b2"] = worst_bunch_b2

    def compute_collision_from_scheme(self) -> tuple[int, int, int]:
        """
        This method reads a filling scheme from a JSON file specified in the configuration, converts
        the filling scheme into boolean arrays for two beams, and calculates the number of
        collisions at IP1 & IP5, IP2, and IP8 by performing convolutions on the arrays.

        Returns:
            tuple[int, int, int]: A tuple containing the number of collisions at IP1 & IP5, IP2, and
                IP8 respectively.

        Raises:
            ValueError: If the filling scheme file is not in JSON format.
            AssertionError: If the length of the beam arrays is not 3564.
        """
        # Get the filling scheme path (in json or csv format)
        filling_scheme_path = self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"]

        # Load the filling scheme
        if not filling_scheme_path.endswith(".json"):
            raise ValueError(
                f"Unknown filling scheme file format: {filling_scheme_path}. It you provided a csv"
                " file, it should have been automatically convert when running the script"
                " 001_make_folders.py. Something went wrong."
            )

        with open(filling_scheme_path, "r") as fid:
            filling_scheme = json.load(fid)

        # Extract booleans beam arrays
        array_b1 = np.array(filling_scheme["beam1"])
        array_b2 = np.array(filling_scheme["beam2"])

        # Assert that the arrays have the required length, and do the convolution
        assert len(array_b1) == len(array_b2) == 3564
        n_collisions_ip1_and_5 = array_b1 @ array_b2
        n_collisions_ip2 = np.roll(array_b1, 891) @ array_b2
        n_collisions_ip8 = np.roll(array_b1, 2670) @ array_b2

        return int(n_collisions_ip1_and_5), int(n_collisions_ip2), int(n_collisions_ip8)

    @property
    def crab(self) -> bool:
        """
        This method checks the configuration settings for the presence and value of the
        "on_crab1" knob. If the knob is present and its value is non-zero, it sets the
        `_crab` attribute to True, indicating that crab cavities are active. Otherwise,
        it sets `_crab` to False.

        Returns:
            bool: True if crab cavities are active, False otherwise.
        """
        if self._crab is None:
            # Get crab cavities
            self._crab = False
            if "on_crab1" in self.config_knobs_and_tuning["knob_settings"]:
                crab_val = float(self.config_knobs_and_tuning["knob_settings"]["on_crab1"])
                if abs(crab_val) > 0:
                    self._crab = True
        return self._crab

    def level_all_by_separation(
        self,
        n_collisions_ip1_and_5: int,
        n_collisions_ip2: int,
        n_collisions_ip8: int,
        collider: xt.Multiline,
    ) -> None:
        """
        This method updates the number of colliding bunches for IP1, IP2, IP5, and IP8 in the
        configuration file and performs luminosity leveling using the provided collider object.
        It also updates the separation knobs for the collider based on the new configuration.

        Args:
            n_collisions_ip1_and_5 (int): Number of collisions at interaction points 1 and 5.
            n_collisions_ip2 (int): Number of collisions at interaction point 2.
            n_collisions_ip8 (int): Number of collisions at interaction point 8.
            collider (xt.Multiline): The collider object to be used for luminosity leveling.

        Returns:
            None
        """
        # Update the number of bunches in the configuration file
        l_n_collisions = [
            n_collisions_ip1_and_5,
            n_collisions_ip2,
            n_collisions_ip1_and_5,
            n_collisions_ip8,
        ]
        for ip, n_collisions in zip(["ip1", "ip2", "ip5", "ip8"], l_n_collisions):
            if ip in self.config_lumi_leveling:
                self.config_lumi_leveling[ip]["num_colliding_bunches"] = n_collisions
            else:
                logging.warning(f"IP {ip} is not in the configuration")

        # ! Crabs are not handled in the following function
        xm.lhc.luminosity_leveling(  # type: ignore
            collider,
            config_lumi_leveling=self.config_lumi_leveling,
            config_beambeam=self.config_beambeam,
        )

        # Update configuration
        if "ip1" in self.config_lumi_leveling:
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip1"], "on_sep1")
        if "ip2" in self.config_lumi_leveling:
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2")
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2h")
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2v")
        if "ip5" in self.config_lumi_leveling:
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip5"], "on_sep5")
        if "ip8" in self.config_lumi_leveling:
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8")
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8h")
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8v")

    def level_ip1_5_by_bunch_intensity(
        self,
        collider: xt.Multiline,
        n_collisions_ip1_and_5: int,
    ) -> None:
        """
        This method modifies the bunch intensity to achieve the desired luminosity
        levels in IP 1 and 5. It updates the configuration with the new intensity values.

        Args:
            collider (xt.Multiline): The collider object containing the beam and lattice
                configuration.
            n_collisions_ip1_and_5 (int):
                The number of collisions in IP 1 and 5.

        Returns:
            None
        """
        # Initial intensity
        bunch_intensity = self.config_beambeam["num_particles_per_bunch"]

        # First level luminosity in IP 1/5 changing the intensity
        if (
            self.config_lumi_leveling_ip1_5 is not None
            and not self.config_lumi_leveling_ip1_5["skip_leveling"]
        ):
            logging.info("Leveling luminosity in IP 1/5 varying the intensity")
            # Update the number of bunches in the configuration file
            self.config_lumi_leveling_ip1_5["num_colliding_bunches"] = n_collisions_ip1_and_5

            # Do the levelling
            bunch_intensity = luminosity_leveling_ip1_5(
                collider,
                self.config_lumi_leveling_ip1_5,
                self.config_beambeam,
                crab=self.crab,
                cross_section=self.config_beambeam["cross_section"],
            )

        # Update the configuration
        self.config_beambeam["final_num_particles_per_bunch"] = float(bunch_intensity)

    def level_ip2_8_by_separation(
        self,
        n_collisions_ip2: int,
        n_collisions_ip8: int,
        collider: xt.Multiline,
    ) -> None:
        """
        This method updates the number of colliding bunches for IP2 and IP8 in the configuration
        file, performs luminosity leveling for the specified collider, and updates the separation
        knobs for both interaction points.

        Args:
            n_collisions_ip2 (int): The number of collisions at interaction point 2 (IP2).
            n_collisions_ip8 (int): The number of collisions at interaction point 8 (IP8).
            collider (xt.Multiline): The collider object for which the luminosity leveling is to be
                performed.

        Returns:
            None
        """
        # Update the number of bunches in the configuration file
        if "ip2" in self.config_lumi_leveling:
            self.config_lumi_leveling["ip2"]["num_colliding_bunches"] = n_collisions_ip2
        if "ip8" in self.config_lumi_leveling:
            self.config_lumi_leveling["ip8"]["num_colliding_bunches"] = n_collisions_ip8

        # Ensure the the num particles per bunch corresponds to the final one
        temp_num_particles_per_bunch = self.config_beambeam["num_particles_per_bunch"]
        if "final_num_particles_per_bunch" in self.config_beambeam:
            self.config_beambeam["num_particles_per_bunch"] = self.config_beambeam[
                "final_num_particles_per_bunch"
            ]
        # Do levelling in IP2 and IP8
        xm.lhc.luminosity_leveling(  # type: ignore
            collider,
            config_lumi_leveling=self.config_lumi_leveling,
            config_beambeam=self.config_beambeam,
        )

        # Update configuration
        if "ip2" in self.config_lumi_leveling:
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2")
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2h")
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2v")
        if "ip8" in self.config_lumi_leveling:
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8")
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8h")
            self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8v")

        # Set back the num particles per bunch to its initial value
        self.config_beambeam["num_particles_per_bunch"] = temp_num_particles_per_bunch

    def add_linear_coupling(self, collider: xt.Multiline) -> None:
        """
        Adds linear coupling to the collider based on the version of the LHC run or HL-LHC optics.

        This method adjusts the collider variables to introduce linear coupling. The specific
        adjustments depend on the version of the LHC run or HL-LHC optics being used.

        Args:
            collider (xt.Multiline): The collider object to which linear coupling will be added.

        Returns:
            None

        Raises:
            ValueError: If the version of the optics or run is unknown.

        Notes:
            - For LHC Run 3.0, the `cmrs.b1_sq` and `cmrs.b2_sq` variables are adjusted.
            - For HL-LHC optics versions 1.6, 1.5, 1.4, and 1.3, the `c_minus_re_b1` and
            `c_minus_re_b2` variables are adjusted.
        """
        # Add linear coupling as the target in the tuning of the base collider was 0
        # (not possible to set it the target to 0.001 for now)
        if self.ver_lhc_run == 3.0:
            collider.vars["cmrs.b1_sq"] += self.config_knobs_and_tuning["delta_cmr"]
            collider.vars["cmrs.b2_sq"] += self.config_knobs_and_tuning["delta_cmr"]
        elif self.ver_hllhc_optics in [1.6, 1.5, 1.4, 1.3]:
            collider.vars["c_minus_re_b1"] += self.config_knobs_and_tuning["delta_cmr"]
            collider.vars["c_minus_re_b2"] += self.config_knobs_and_tuning["delta_cmr"]
        else:
            raise ValueError(
                f"Unknown version of the optics/run: {self.ver_hllhc_optics}, {self.ver_lhc_run}."
            )

    def assert_tune_chroma_coupling(self, collider: xt.Multiline) -> None:
        """
        Asserts that the tune, chromaticity, and linear coupling of the collider
        match the expected values specified in the configuration.

        Args:
            collider (xt.Multiline): The collider object containing the lines to be checked.

        Returns:
            None

        Raises:
            AssertionError: If any of the tune, chromaticity, or linear coupling values do not match
                the expected values within the specified tolerances.

        Notes:
            The function checks the following parameters for each line ("lhcb1" and "lhcb2"):
            - Horizontal tune (qx)
            - Vertical tune (qy)
            - Horizontal chromaticity (dqx)
            - Vertical chromaticity (dqy)
            - Linear coupling (c_minus)

        The expected values are retrieved from the `self.config_knobs_and_tuning` dictionary.
        """
        for line_name in ["lhcb1", "lhcb2"]:
            tw = collider[line_name].twiss()
            assert np.isclose(tw.qx, self.config_knobs_and_tuning["qx"][line_name], atol=1e-4), (
                f"tune_x is not correct for {line_name}. Expected"
                f" {self.config_knobs_and_tuning['qx'][line_name]}, got {tw.qx}"
            )
            assert np.isclose(tw.qy, self.config_knobs_and_tuning["qy"][line_name], atol=1e-4), (
                f"tune_y is not correct for {line_name}. Expected"
                f" {self.config_knobs_and_tuning['qy'][line_name]}, got {tw.qy}"
            )
            assert np.isclose(
                tw.dqx,
                self.config_knobs_and_tuning["dqx"][line_name],
                rtol=1e-2,
            ), (
                f"chromaticity_x is not correct for {line_name}. Expected"
                f" {self.config_knobs_and_tuning['dqx'][line_name]}, got {tw.dqx}"
            )
            assert np.isclose(
                tw.dqy,
                self.config_knobs_and_tuning["dqy"][line_name],
                rtol=1e-2,
            ), (
                f"chromaticity_y is not correct for {line_name}. Expected"
                f" {self.config_knobs_and_tuning['dqy'][line_name]}, got {tw.dqy}"
            )

            assert np.isclose(
                tw.c_minus,
                self.config_knobs_and_tuning["delta_cmr"],
                atol=5e-3,
            ), (
                f"linear coupling is not correct for {line_name}. Expected"
                f" {self.config_knobs_and_tuning['delta_cmr']}, got {tw.c_minus}"
            )

    def record_beta_functions(self, collider: xt.Multiline) -> None:
        """
        Records the beta functions at the IPs in the collider.

        Args:
            collider (xt.Multiline): The collider object to record the beta functions.

        Returns:
            None
        """
        # Record beta functions at the IPs
        for ip in ["ip1", "ip2", "ip5", "ip8"]:
            tw = collider.lhcb1.twiss()
            self.config_collider[f"beta_x_{ip}"] = float(np.round(float(tw["betx", ip]), 5))
            self.config_collider[f"beta_y_{ip}"] = float(np.round(float(tw["bety", ip]), 5))

    def configure_beam_beam(self, collider: xt.Multiline) -> None:
        """
        Configures the beam-beam interactions for the collider.

        This method sets up the beam-beam interactions by configuring the number of particles per
        bunch, the horizontal emittance (nemitt_x), and the vertical emittance (nemitt_y) based on
        the provided configuration. Additionally, it configures the filling scheme mask and bunch
        numbers if a filling pattern is specified in the configuration.

        Args:
            collider (xt.Multiline): The collider object to configure.

        Returns:
            None
        """
        collider.configure_beambeam_interactions(
            num_particles=self.config_beambeam["num_particles_per_bunch"],
            nemitt_x=self.config_beambeam["nemitt_x"],
            nemitt_y=self.config_beambeam["nemitt_y"],
        )

        # Configure filling scheme mask and bunch numbers
        if "mask_with_filling_pattern" in self.config_beambeam and (
            "pattern_fname" in self.config_beambeam["mask_with_filling_pattern"]
            and self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"] is not None
        ):
            fname = self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"]
            with open(fname, "r") as fid:
                filling = json.load(fid)
            filling_pattern_cw = filling["beam1"]
            filling_pattern_acw = filling["beam2"]

            # Initialize bunch numbers with empty values
            i_bunch_cw = None
            i_bunch_acw = None

            # Only track bunch number if a filling pattern has been provided
            if "i_bunch_b1" in self.config_beambeam["mask_with_filling_pattern"]:
                i_bunch_cw = self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"]
            if "i_bunch_b2" in self.config_beambeam["mask_with_filling_pattern"]:
                i_bunch_acw = self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b2"]

            # Note that a bunch number must be provided if a filling pattern is provided
            # Apply filling pattern
            collider.apply_filling_pattern(
                filling_pattern_cw=filling_pattern_cw,
                filling_pattern_acw=filling_pattern_acw,
                i_bunch_cw=i_bunch_cw,
                i_bunch_acw=i_bunch_acw,
            )

    def record_final_luminosity(self, collider: xt.Multiline, l_n_collisions: list[int]) -> None:
        """
        Records the final luminosity and pile-up for specified interaction points (IPs)
        in the collider, both with and without beam-beam effects.

        Args:
            collider : (xt.Multiline): The collider object configured.
            l_n_collisions (list[int]): A list containing the number of colliding bunches for each
                IP.

        Returns:
            None
        """
        # Define IPs in which the luminosity will be computed
        l_ip = ["ip1", "ip2", "ip5", "ip8"]

        # Ensure that the final number of particles per bunch is defined, even
        # if the leveling has been done by separation
        if "final_num_particles_per_bunch" not in self.config_beambeam:
            self.config_beambeam["final_num_particles_per_bunch"] = self.config_beambeam[
                "num_particles_per_bunch"
            ]

        def _twiss_and_compute_lumi(collider, l_n_collisions):
            # Loop over each IP and record the luminosity
            twiss_b1 = collider["lhcb1"].twiss()
            twiss_b2 = collider["lhcb2"].twiss()
            l_lumi = []
            l_PU = []
            for n_col, ip in zip(l_n_collisions, l_ip):
                L = xt.lumi.luminosity_from_twiss(  # type: ignore
                    n_colliding_bunches=n_col,
                    num_particles_per_bunch=self.config_beambeam["final_num_particles_per_bunch"],
                    ip_name=ip,
                    nemitt_x=self.config_beambeam["nemitt_x"],
                    nemitt_y=self.config_beambeam["nemitt_y"],
                    sigma_z=self.config_beambeam["sigma_z"],
                    twiss_b1=twiss_b1,
                    twiss_b2=twiss_b2,
                    crab=self.crab,
                )
                PU = compute_PU(
                    L,
                    n_col,
                    twiss_b1["T_rev0"],
                    cross_section=self.config_beambeam["cross_section"],
                )

                l_lumi.append(L)
                l_PU.append(PU)

            return l_lumi, l_PU

        # Get the final luminosity in all IPs, without beam-beam
        collider.vars["beambeam_scale"] = 0
        l_lumi, l_PU = _twiss_and_compute_lumi(collider, l_n_collisions)

        # Update configuration
        for ip, L, PU in zip(l_ip, l_lumi, l_PU):
            self.config_beambeam[f"luminosity_{ip}_without_beam_beam"] = float(L)
            self.config_beambeam[f"Pile-up_{ip}_without_beam_beam"] = float(PU)

        # Get the final luminosity in all IPs, with beam-beam
        collider.vars["beambeam_scale"] = 1
        l_lumi, l_PU = _twiss_and_compute_lumi(collider, l_n_collisions)

        # Update configuration
        for ip, L, PU in zip(l_ip, l_lumi, l_PU):
            self.config_beambeam[f"luminosity_{ip}_with_beam_beam"] = float(L)
            self.config_beambeam[f"Pile-up_{ip}_with_beam_beam"] = float(PU)

    def write_collider_to_disk(self, collider, full_configuration) -> None:
        """
        Writes the collider object to disk in JSON format if the save_output_collider flag is set.

        Args:
            collider (Collider): The collider object to be saved.
            full_configuration (dict): The full configuration dictionary to be deep-copied into the
                collider's metadata.

        Returns:
            None
        """
        if self.save_output_collider:
            logging.info("Saving collider as json")
            if (
                hasattr(collider, "metadata")
                and collider.metadata is not None
                and isinstance(collider.metadata, dict)
            ):
                collider.metadata.update(copy.deepcopy(full_configuration))
            else:
                collider.metadata = copy.deepcopy(full_configuration)
            collider.to_json(self.path_collider_file_for_tracking_as_output)

            # Compress the collider file to zip to ease the load on afs
            if self.compress:
                compress_and_write(self.path_collider_file_for_tracking_as_output)

    @staticmethod
    def update_configuration_knob(
        collider: xt.Multiline, dictionnary: dict, knob_name: str
    ) -> None:
        """
        Updates the given dictionary with the final value of a specified knob from the collider.

        Args:
            collider (xt.Multiline): The collider object containing various variables.
            dictionnary (dict): The dictionary to be updated with the knob's final value.
            knob_name (str): The name of the knob whose value is to be retrieved and stored.

        Returns:
            None
        """
        if knob_name in collider.vars.keys():
            dictionnary[f"final_{knob_name}"] = float(collider.vars[knob_name]._value)
        else:
            logging.warning(f"Knob {knob_name} not found in the collider")

    @staticmethod
    def return_fingerprint(collider, line_name="lhcb1") -> str:
        """
        Generate a detailed fingerprint of the specified collider line. Useful to compare two
        colliders.

        Args:
            collider (xt.Multiline): The collider object containing the line data.
            line_name (str): The name of the line to analyze within the collider. Default to "lhcb1".

        Returns:
            str:
                A formatted string containing detailed information about the collider line, including:
                - Installed element types
                - Tunes and chromaticity
                - Synchrotron tune and slip factor
                - Twiss parameters and phases at interaction points (IPs)
                - Dispersion and crab dispersion at IPs
                - Amplitude detuning coefficients
                - Non-linear chromaticity
                - Tunes and momentum compaction vs delta
        """
        line = collider[line_name]

        tw = line.twiss()
        tt = line.get_table()

        det = line.get_amplitude_detuning_coefficients(a0_sigmas=0.1, a1_sigmas=0.2, a2_sigmas=0.3)

        det_table = xt.Table(
            {
                "name": np.array(list(det.keys())),
                "value": np.array(list(det.values())),
            }
        )

        nl_chrom = line.get_non_linear_chromaticity(
            delta0_range=(-2e-4, 2e-4), num_delta=5, fit_order=3
        )

        out = ""

        out += f"Line: {line_name}\n"
        out += "\n"

        out += "Installed element types:\n"
        out += repr([nn for nn in sorted(list(set(tt.element_type))) if len(nn) > 0]) + "\n"
        out += "\n"

        out += f'Tunes:        Qx  = {tw["qx"]:.5f}       Qy = {tw["qy"]:.5f}\n'
        out += f"""Chromaticity: Q'x = {tw["dqx"]:.2f}     Q'y = """ + f'{tw["dqy"]:.2f}\n'
        out += f'c_minus:      {tw["c_minus"]:.5e}\n'
        out += "\n"

        out += f'Synchrotron tune: {tw["qs"]:5e}\n'
        out += f'Slip factor:      {tw["slip_factor"]:.5e}\n'
        out += "\n"

        out += "Twiss parameters and phases at IPs:\n"
        out += (
            tw.rows["ip.*"]
            .cols["name s betx bety alfx alfy mux muy"]
            .show(output=str, max_col_width=int(1e6), digits=8)
        )
        out += "\n\n"

        out += "Dispersion at IPs:\n"
        out += (
            tw.rows["ip.*"]
            .cols["name s dx dy dpx dpy"]
            .show(output=str, max_col_width=int(1e6), digits=8)
        )
        out += "\n\n"

        out += "Crab dispersion at IPs:\n"
        out += (
            tw.rows["ip.*"]
            .cols["name s dx_zeta dy_zeta dpx_zeta dpy_zeta"]
            .show(output=str, max_col_width=int(1e6), digits=8)
        )
        out += "\n\n"

        out += "Amplitude detuning coefficients:\n"
        out += det_table.show(output=str, max_col_width=int(1e6), digits=6)
        out += "\n\n"

        out += "Non-linear chromaticity:\n"
        out += f'dnqx = {list(nl_chrom["dnqx"])}\n'
        out += f'dnqy = {list(nl_chrom["dnqy"])}\n'
        out += "\n\n"

        out += "Tunes and momentum compaction vs delta:\n"
        out += nl_chrom.show(output=str, max_col_width=int(1e6), digits=6)
        out += "\n\n"

        return out

crab: bool property

This method checks the configuration settings for the presence and value of the "on_crab1" knob. If the knob is present and its value is non-zero, it sets the _crab attribute to True, indicating that crab cavities are active. Otherwise, it sets _crab to False.

Returns:

Name Type Description
bool bool

True if crab cavities are active, False otherwise.

dict_orbit_correction: dict property

Generates and returns a dictionary containing orbit correction parameters.

This method checks if the orbit correction dictionary has already been generated. If not, it determines the appropriate set of orbit correction parameters based on the version of HLLHC optics or LHC run provided.

Returns:

Name Type Description
dict dict

A dictionary containing orbit correction parameters.

Raises:

Type Description
ValueError

If both ver_hllhc_optics and ver_lhc_run are defined.
ValueError

If no optics specific tools are available for the provided configuration.

__init__(configuration, path_collider_file_for_configuration_as_input, ver_hllhc_optics, ver_lhc_run, ions)

Initialize the XsuiteCollider class with the given configuration and parameters.

Parameters:

Name Type Description Default
configuration dict

A dictionary containing various configuration settings. - config_beambeam (dict): Configuration for beam-beam interactions. - config_knobs_and_tuning (dict): Configuration for knobs and tuning. - config_lumi_leveling (dict): Configuration for luminosity leveling. - save_output_collider (bool): Flag to save the final collider to disk. - path_collider_file_for_tracking_as_output (str): Path to save the final collider. - config_lumi_leveling_ip1_5 (optional): Configuration for luminosity leveling at IP1 and IP5.

 required
path_collider_file_for_configuration_as_input str

Path to the collider file.

 required
ver_hllhc_optics float

Version of the HL-LHC optics.

 required
ver_lhc_run float

Version of the LHC run.

 required
ions bool

Flag indicating if ions are used.

 required
Source code in study_da/generate/master_classes/xsuite_collider.py 
 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
def __init__(
    self,
    configuration: dict,
    path_collider_file_for_configuration_as_input: str,
    ver_hllhc_optics: float,
    ver_lhc_run: float,
    ions: bool,
):
    """
    Initialize the XsuiteCollider class with the given configuration and parameters.

    Args:
        configuration (dict): A dictionary containing various configuration settings.
            - config_beambeam (dict): Configuration for beam-beam interactions.
            - config_knobs_and_tuning (dict): Configuration for knobs and tuning.
            - config_lumi_leveling (dict): Configuration for luminosity leveling.
            - save_output_collider (bool): Flag to save the final collider to disk.
            - path_collider_file_for_tracking_as_output (str): Path to save the final collider.
            - config_lumi_leveling_ip1_5 (optional): Configuration for luminosity leveling at
                IP1 and IP5.
        path_collider_file_for_configuration_as_input (str): Path to the collider file.
        ver_hllhc_optics (float): Version of the HL-LHC optics.
        ver_lhc_run (float): Version of the LHC run.
        ions (bool): Flag indicating if ions are used.
    """
    # Collider file path
    self.path_collider_file_for_configuration_as_input = (
        path_collider_file_for_configuration_as_input
    )

    # Configuration variables
    self.config_beambeam: dict[str, Any] = configuration["config_beambeam"]
    self.config_knobs_and_tuning: dict[str, Any] = configuration["config_knobs_and_tuning"]
    self.config_lumi_leveling: dict[str, Any] = configuration["config_lumi_leveling"]

    # self.config_lumi_leveling_ip1_5 will be None if not present in the configuration
    self.config_lumi_leveling_ip1_5: dict[str, Any] = configuration.get(
        "config_lumi_leveling_ip1_5"
    )

    # Collider configuration
    self.config_collider: dict[str, Any] = configuration

    # Optics version (needed to select the appropriate optics specific functions)
    self.ver_hllhc_optics: float = ver_hllhc_optics
    self.ver_lhc_run: float = ver_lhc_run
    self.ions: bool = ions
    self._dict_orbit_correction: dict | None = None

    # Crab cavities
    self._crab: bool | None = None

    # Save collider to disk
    self.save_output_collider = configuration["save_output_collider"]
    self.path_collider_file_for_tracking_as_output = configuration[
        "path_collider_file_for_tracking_as_output"
    ]
    self.compress = configuration["compress"]

add_linear_coupling(collider)

Adds linear coupling to the collider based on the version of the LHC run or HL-LHC optics.

This method adjusts the collider variables to introduce linear coupling. The specific adjustments depend on the version of the LHC run or HL-LHC optics being used.

Parameters:

Name Type Description Default
collider Multiline

The collider object to which linear coupling will be added.

 required

Returns:

Type Description
None

None

Raises:

Type Description
ValueError

If the version of the optics or run is unknown.
Notes
  • For LHC Run 3.0, the cmrs.b1_sq and cmrs.b2_sq variables are adjusted.
  • For HL-LHC optics versions 1.6, 1.5, 1.4, and 1.3, the c_minus_re_b1 and c_minus_re_b2 variables are adjusted.
Source code in study_da/generate/master_classes/xsuite_collider.py 
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
def add_linear_coupling(self, collider: xt.Multiline) -> None:
    """
    Adds linear coupling to the collider based on the version of the LHC run or HL-LHC optics.

    This method adjusts the collider variables to introduce linear coupling. The specific
    adjustments depend on the version of the LHC run or HL-LHC optics being used.

    Args:
        collider (xt.Multiline): The collider object to which linear coupling will be added.

    Returns:
        None

    Raises:
        ValueError: If the version of the optics or run is unknown.

    Notes:
        - For LHC Run 3.0, the `cmrs.b1_sq` and `cmrs.b2_sq` variables are adjusted.
        - For HL-LHC optics versions 1.6, 1.5, 1.4, and 1.3, the `c_minus_re_b1` and
        `c_minus_re_b2` variables are adjusted.
    """
    # Add linear coupling as the target in the tuning of the base collider was 0
    # (not possible to set it the target to 0.001 for now)
    if self.ver_lhc_run == 3.0:
        collider.vars["cmrs.b1_sq"] += self.config_knobs_and_tuning["delta_cmr"]
        collider.vars["cmrs.b2_sq"] += self.config_knobs_and_tuning["delta_cmr"]
    elif self.ver_hllhc_optics in [1.6, 1.5, 1.4, 1.3]:
        collider.vars["c_minus_re_b1"] += self.config_knobs_and_tuning["delta_cmr"]
        collider.vars["c_minus_re_b2"] += self.config_knobs_and_tuning["delta_cmr"]
    else:
        raise ValueError(
            f"Unknown version of the optics/run: {self.ver_hllhc_optics}, {self.ver_lhc_run}."
        )

assert_tune_chroma_coupling(collider)

Asserts that the tune, chromaticity, and linear coupling of the collider match the expected values specified in the configuration.

Parameters:

Name Type Description Default
collider Multiline

The collider object containing the lines to be checked.

 required

Returns:

Type Description
None

None

Raises:

Type Description
AssertionError

If any of the tune, chromaticity, or linear coupling values do not match the expected values within the specified tolerances.
Notes

The function checks the following parameters for each line ("lhcb1" and "lhcb2"): - Horizontal tune (qx) - Vertical tune (qy) - Horizontal chromaticity (dqx) - Vertical chromaticity (dqy) - Linear coupling (c_minus)

The expected values are retrieved from the self.config_knobs_and_tuning dictionary.

Source code in study_da/generate/master_classes/xsuite_collider.py 
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
def assert_tune_chroma_coupling(self, collider: xt.Multiline) -> None:
    """
    Asserts that the tune, chromaticity, and linear coupling of the collider
    match the expected values specified in the configuration.

    Args:
        collider (xt.Multiline): The collider object containing the lines to be checked.

    Returns:
        None

    Raises:
        AssertionError: If any of the tune, chromaticity, or linear coupling values do not match
            the expected values within the specified tolerances.

    Notes:
        The function checks the following parameters for each line ("lhcb1" and "lhcb2"):
        - Horizontal tune (qx)
        - Vertical tune (qy)
        - Horizontal chromaticity (dqx)
        - Vertical chromaticity (dqy)
        - Linear coupling (c_minus)

    The expected values are retrieved from the `self.config_knobs_and_tuning` dictionary.
    """
    for line_name in ["lhcb1", "lhcb2"]:
        tw = collider[line_name].twiss()
        assert np.isclose(tw.qx, self.config_knobs_and_tuning["qx"][line_name], atol=1e-4), (
            f"tune_x is not correct for {line_name}. Expected"
            f" {self.config_knobs_and_tuning['qx'][line_name]}, got {tw.qx}"
        )
        assert np.isclose(tw.qy, self.config_knobs_and_tuning["qy"][line_name], atol=1e-4), (
            f"tune_y is not correct for {line_name}. Expected"
            f" {self.config_knobs_and_tuning['qy'][line_name]}, got {tw.qy}"
        )
        assert np.isclose(
            tw.dqx,
            self.config_knobs_and_tuning["dqx"][line_name],
            rtol=1e-2,
        ), (
            f"chromaticity_x is not correct for {line_name}. Expected"
            f" {self.config_knobs_and_tuning['dqx'][line_name]}, got {tw.dqx}"
        )
        assert np.isclose(
            tw.dqy,
            self.config_knobs_and_tuning["dqy"][line_name],
            rtol=1e-2,
        ), (
            f"chromaticity_y is not correct for {line_name}. Expected"
            f" {self.config_knobs_and_tuning['dqy'][line_name]}, got {tw.dqy}"
        )

        assert np.isclose(
            tw.c_minus,
            self.config_knobs_and_tuning["delta_cmr"],
            atol=5e-3,
        ), (
            f"linear coupling is not correct for {line_name}. Expected"
            f" {self.config_knobs_and_tuning['delta_cmr']}, got {tw.c_minus}"
        )

compute_collision_from_scheme()

This method reads a filling scheme from a JSON file specified in the configuration, converts the filling scheme into boolean arrays for two beams, and calculates the number of collisions at IP1 & IP5, IP2, and IP8 by performing convolutions on the arrays.

Returns:

Type Description
tuple[int, int, int]

tuple[int, int, int]: A tuple containing the number of collisions at IP1 & IP5, IP2, and IP8 respectively.

Raises:

Type Description
ValueError

If the filling scheme file is not in JSON format.
AssertionError

If the length of the beam arrays is not 3564.
Source code in study_da/generate/master_classes/xsuite_collider.py 
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
def compute_collision_from_scheme(self) -> tuple[int, int, int]:
    """
    This method reads a filling scheme from a JSON file specified in the configuration, converts
    the filling scheme into boolean arrays for two beams, and calculates the number of
    collisions at IP1 & IP5, IP2, and IP8 by performing convolutions on the arrays.

    Returns:
        tuple[int, int, int]: A tuple containing the number of collisions at IP1 & IP5, IP2, and
            IP8 respectively.

    Raises:
        ValueError: If the filling scheme file is not in JSON format.
        AssertionError: If the length of the beam arrays is not 3564.
    """
    # Get the filling scheme path (in json or csv format)
    filling_scheme_path = self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"]

    # Load the filling scheme
    if not filling_scheme_path.endswith(".json"):
        raise ValueError(
            f"Unknown filling scheme file format: {filling_scheme_path}. It you provided a csv"
            " file, it should have been automatically convert when running the script"
            " 001_make_folders.py. Something went wrong."
        )

    with open(filling_scheme_path, "r") as fid:
        filling_scheme = json.load(fid)

    # Extract booleans beam arrays
    array_b1 = np.array(filling_scheme["beam1"])
    array_b2 = np.array(filling_scheme["beam2"])

    # Assert that the arrays have the required length, and do the convolution
    assert len(array_b1) == len(array_b2) == 3564
    n_collisions_ip1_and_5 = array_b1 @ array_b2
    n_collisions_ip2 = np.roll(array_b1, 891) @ array_b2
    n_collisions_ip8 = np.roll(array_b1, 2670) @ array_b2

    return int(n_collisions_ip1_and_5), int(n_collisions_ip2), int(n_collisions_ip8)

configure_beam_beam(collider)

Configures the beam-beam interactions for the collider.

This method sets up the beam-beam interactions by configuring the number of particles per bunch, the horizontal emittance (nemitt_x), and the vertical emittance (nemitt_y) based on the provided configuration. Additionally, it configures the filling scheme mask and bunch numbers if a filling pattern is specified in the configuration.

Parameters:

Name Type Description Default
collider Multiline

The collider object to configure.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
def configure_beam_beam(self, collider: xt.Multiline) -> None:
    """
    Configures the beam-beam interactions for the collider.

    This method sets up the beam-beam interactions by configuring the number of particles per
    bunch, the horizontal emittance (nemitt_x), and the vertical emittance (nemitt_y) based on
    the provided configuration. Additionally, it configures the filling scheme mask and bunch
    numbers if a filling pattern is specified in the configuration.

    Args:
        collider (xt.Multiline): The collider object to configure.

    Returns:
        None
    """
    collider.configure_beambeam_interactions(
        num_particles=self.config_beambeam["num_particles_per_bunch"],
        nemitt_x=self.config_beambeam["nemitt_x"],
        nemitt_y=self.config_beambeam["nemitt_y"],
    )

    # Configure filling scheme mask and bunch numbers
    if "mask_with_filling_pattern" in self.config_beambeam and (
        "pattern_fname" in self.config_beambeam["mask_with_filling_pattern"]
        and self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"] is not None
    ):
        fname = self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"]
        with open(fname, "r") as fid:
            filling = json.load(fid)
        filling_pattern_cw = filling["beam1"]
        filling_pattern_acw = filling["beam2"]

        # Initialize bunch numbers with empty values
        i_bunch_cw = None
        i_bunch_acw = None

        # Only track bunch number if a filling pattern has been provided
        if "i_bunch_b1" in self.config_beambeam["mask_with_filling_pattern"]:
            i_bunch_cw = self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"]
        if "i_bunch_b2" in self.config_beambeam["mask_with_filling_pattern"]:
            i_bunch_acw = self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b2"]

        # Note that a bunch number must be provided if a filling pattern is provided
        # Apply filling pattern
        collider.apply_filling_pattern(
            filling_pattern_cw=filling_pattern_cw,
            filling_pattern_acw=filling_pattern_acw,
            i_bunch_cw=i_bunch_cw,
            i_bunch_acw=i_bunch_acw,
        )

install_beam_beam_wrapper(collider)

This method installs beam-beam interactions in the collider with the specified parameters. The beam-beam lenses are initially inactive and not configured.

Parameters:

Name Type Description Default
collider Multiline

The collider object where the beam-beam interactions will be installed.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
def install_beam_beam_wrapper(self, collider: xt.Multiline) -> None:
    """
    This method installs beam-beam interactions in the collider with the specified
    parameters. The beam-beam lenses are initially inactive and not configured.

    Args:
        collider (xt.Multiline): The collider object where the beam-beam interactions
            will be installed.

    Returns:
        None
    """
    # Install beam-beam lenses (inactive and not configured)
    collider.install_beambeam_interactions(
        clockwise_line="lhcb1",
        anticlockwise_line="lhcb2",
        ip_names=["ip1", "ip2", "ip5", "ip8"],
        delay_at_ips_slots=[0, 891, 0, 2670],
        num_long_range_encounters_per_side=self.config_beambeam[
            "num_long_range_encounters_per_side"
        ],
        num_slices_head_on=self.config_beambeam["num_slices_head_on"],
        harmonic_number=35640,
        bunch_spacing_buckets=self.config_beambeam["bunch_spacing_buckets"],
        sigmaz=self.config_beambeam["sigma_z"],
    )

level_all_by_separation(n_collisions_ip1_and_5, n_collisions_ip2, n_collisions_ip8, collider)

This method updates the number of colliding bunches for IP1, IP2, IP5, and IP8 in the configuration file and performs luminosity leveling using the provided collider object. It also updates the separation knobs for the collider based on the new configuration.

Parameters:

Name Type Description Default
n_collisions_ip1_and_5 int

Number of collisions at interaction points 1 and 5.

 required
n_collisions_ip2 int

Number of collisions at interaction point 2.

 required
n_collisions_ip8 int

Number of collisions at interaction point 8.

 required
collider Multiline

The collider object to be used for luminosity leveling.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
def level_all_by_separation(
    self,
    n_collisions_ip1_and_5: int,
    n_collisions_ip2: int,
    n_collisions_ip8: int,
    collider: xt.Multiline,
) -> None:
    """
    This method updates the number of colliding bunches for IP1, IP2, IP5, and IP8 in the
    configuration file and performs luminosity leveling using the provided collider object.
    It also updates the separation knobs for the collider based on the new configuration.

    Args:
        n_collisions_ip1_and_5 (int): Number of collisions at interaction points 1 and 5.
        n_collisions_ip2 (int): Number of collisions at interaction point 2.
        n_collisions_ip8 (int): Number of collisions at interaction point 8.
        collider (xt.Multiline): The collider object to be used for luminosity leveling.

    Returns:
        None
    """
    # Update the number of bunches in the configuration file
    l_n_collisions = [
        n_collisions_ip1_and_5,
        n_collisions_ip2,
        n_collisions_ip1_and_5,
        n_collisions_ip8,
    ]
    for ip, n_collisions in zip(["ip1", "ip2", "ip5", "ip8"], l_n_collisions):
        if ip in self.config_lumi_leveling:
            self.config_lumi_leveling[ip]["num_colliding_bunches"] = n_collisions
        else:
            logging.warning(f"IP {ip} is not in the configuration")

    # ! Crabs are not handled in the following function
    xm.lhc.luminosity_leveling(  # type: ignore
        collider,
        config_lumi_leveling=self.config_lumi_leveling,
        config_beambeam=self.config_beambeam,
    )

    # Update configuration
    if "ip1" in self.config_lumi_leveling:
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip1"], "on_sep1")
    if "ip2" in self.config_lumi_leveling:
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2")
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2h")
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2v")
    if "ip5" in self.config_lumi_leveling:
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip5"], "on_sep5")
    if "ip8" in self.config_lumi_leveling:
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8")
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8h")
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8v")

level_ip1_5_by_bunch_intensity(collider, n_collisions_ip1_and_5)

This method modifies the bunch intensity to achieve the desired luminosity levels in IP 1 and 5. It updates the configuration with the new intensity values.

Parameters:

Name Type Description Default
collider Multiline

The collider object containing the beam and lattice configuration.

 required
n_collisions_ip1_and_5 int

The number of collisions in IP 1 and 5.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
def level_ip1_5_by_bunch_intensity(
    self,
    collider: xt.Multiline,
    n_collisions_ip1_and_5: int,
) -> None:
    """
    This method modifies the bunch intensity to achieve the desired luminosity
    levels in IP 1 and 5. It updates the configuration with the new intensity values.

    Args:
        collider (xt.Multiline): The collider object containing the beam and lattice
            configuration.
        n_collisions_ip1_and_5 (int):
            The number of collisions in IP 1 and 5.

    Returns:
        None
    """
    # Initial intensity
    bunch_intensity = self.config_beambeam["num_particles_per_bunch"]

    # First level luminosity in IP 1/5 changing the intensity
    if (
        self.config_lumi_leveling_ip1_5 is not None
        and not self.config_lumi_leveling_ip1_5["skip_leveling"]
    ):
        logging.info("Leveling luminosity in IP 1/5 varying the intensity")
        # Update the number of bunches in the configuration file
        self.config_lumi_leveling_ip1_5["num_colliding_bunches"] = n_collisions_ip1_and_5

        # Do the levelling
        bunch_intensity = luminosity_leveling_ip1_5(
            collider,
            self.config_lumi_leveling_ip1_5,
            self.config_beambeam,
            crab=self.crab,
            cross_section=self.config_beambeam["cross_section"],
        )

    # Update the configuration
    self.config_beambeam["final_num_particles_per_bunch"] = float(bunch_intensity)

level_ip2_8_by_separation(n_collisions_ip2, n_collisions_ip8, collider)

This method updates the number of colliding bunches for IP2 and IP8 in the configuration file, performs luminosity leveling for the specified collider, and updates the separation knobs for both interaction points.

Parameters:

Name Type Description Default
n_collisions_ip2 int

The number of collisions at interaction point 2 (IP2).

 required
n_collisions_ip8 int

The number of collisions at interaction point 8 (IP8).

 required
collider Multiline

The collider object for which the luminosity leveling is to be performed.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
def level_ip2_8_by_separation(
    self,
    n_collisions_ip2: int,
    n_collisions_ip8: int,
    collider: xt.Multiline,
) -> None:
    """
    This method updates the number of colliding bunches for IP2 and IP8 in the configuration
    file, performs luminosity leveling for the specified collider, and updates the separation
    knobs for both interaction points.

    Args:
        n_collisions_ip2 (int): The number of collisions at interaction point 2 (IP2).
        n_collisions_ip8 (int): The number of collisions at interaction point 8 (IP8).
        collider (xt.Multiline): The collider object for which the luminosity leveling is to be
            performed.

    Returns:
        None
    """
    # Update the number of bunches in the configuration file
    if "ip2" in self.config_lumi_leveling:
        self.config_lumi_leveling["ip2"]["num_colliding_bunches"] = n_collisions_ip2
    if "ip8" in self.config_lumi_leveling:
        self.config_lumi_leveling["ip8"]["num_colliding_bunches"] = n_collisions_ip8

    # Ensure the the num particles per bunch corresponds to the final one
    temp_num_particles_per_bunch = self.config_beambeam["num_particles_per_bunch"]
    if "final_num_particles_per_bunch" in self.config_beambeam:
        self.config_beambeam["num_particles_per_bunch"] = self.config_beambeam[
            "final_num_particles_per_bunch"
        ]
    # Do levelling in IP2 and IP8
    xm.lhc.luminosity_leveling(  # type: ignore
        collider,
        config_lumi_leveling=self.config_lumi_leveling,
        config_beambeam=self.config_beambeam,
    )

    # Update configuration
    if "ip2" in self.config_lumi_leveling:
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2")
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2h")
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip2"], "on_sep2v")
    if "ip8" in self.config_lumi_leveling:
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8")
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8h")
        self.update_configuration_knob(collider, self.config_lumi_leveling["ip8"], "on_sep8v")

    # Set back the num particles per bunch to its initial value
    self.config_beambeam["num_particles_per_bunch"] = temp_num_particles_per_bunch

load_collider()

Load a collider configuration from a file.

If the file path ends with ".zip", the file is uncompressed locally and the collider configuration is loaded from the uncompressed file. Otherwise, the collider configuration is loaded directly from the file.

Returns:

Type Description
Multiline

xt.Multiline: The loaded collider configuration.
Source code in study_da/generate/master_classes/xsuite_collider.py 
213
214
215
216
217
218
219
220
221
222
223
224
def load_collider(self) -> xt.Multiline:
    """
    Load a collider configuration from a file.

    If the file path ends with ".zip", the file is uncompressed locally
    and the collider configuration is loaded from the uncompressed file.
    Otherwise, the collider configuration is loaded directly from the file.

    Returns:
        xt.Multiline: The loaded collider configuration.
    """
    return self._load_collider(self.path_collider_file_for_configuration_as_input)

match_tune_and_chroma(collider, match_linear_coupling_to_zero=True)

This method adjusts the tune and chromaticity of the specified collider lines ("lhcb1" and "lhcb2") to the target values defined in the configuration. It also optionally matches the linear coupling to zero.

Parameters:

Name Type Description Default
collider Multiline

The collider object containing the lines to be tuned.

 required
match_linear_coupling_to_zero bool

If True, linear coupling will be matched to zero. Defaults to True.

 True

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
def match_tune_and_chroma(
    self, collider: xt.Multiline, match_linear_coupling_to_zero: bool = True
) -> None:
    """
    This method adjusts the tune and chromaticity of the specified collider lines
    ("lhcb1" and "lhcb2") to the target values defined in the configuration. It also
    optionally matches the linear coupling to zero.

    Args:
        collider (xt.Multiline): The collider object containing the lines to be tuned.
        match_linear_coupling_to_zero (bool, optional): If True, linear coupling will be
            matched to zero. Defaults to True.

    Returns:
        None
    """
    for line_name in ["lhcb1", "lhcb2"]:
        knob_names = self.config_knobs_and_tuning["knob_names"][line_name]

        targets = {
            "qx": self.config_knobs_and_tuning["qx"][line_name],
            "qy": self.config_knobs_and_tuning["qy"][line_name],
            "dqx": self.config_knobs_and_tuning["dqx"][line_name],
            "dqy": self.config_knobs_and_tuning["dqy"][line_name],
        }

        xm.machine_tuning(
            line=collider[line_name],
            enable_closed_orbit_correction=True,
            enable_linear_coupling_correction=match_linear_coupling_to_zero,
            enable_tune_correction=True,
            enable_chromaticity_correction=True,
            knob_names=knob_names,
            targets=targets,
            line_co_ref=collider[f"{line_name}_co_ref"],
            co_corr_config=self.dict_orbit_correction[line_name],
        )

record_beta_functions(collider)

Records the beta functions at the IPs in the collider.

Parameters:

Name Type Description Default
collider Multiline

The collider object to record the beta functions.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
def record_beta_functions(self, collider: xt.Multiline) -> None:
    """
    Records the beta functions at the IPs in the collider.

    Args:
        collider (xt.Multiline): The collider object to record the beta functions.

    Returns:
        None
    """
    # Record beta functions at the IPs
    for ip in ["ip1", "ip2", "ip5", "ip8"]:
        tw = collider.lhcb1.twiss()
        self.config_collider[f"beta_x_{ip}"] = float(np.round(float(tw["betx", ip]), 5))
        self.config_collider[f"beta_y_{ip}"] = float(np.round(float(tw["bety", ip]), 5))

record_final_luminosity(collider, l_n_collisions)

Records the final luminosity and pile-up for specified interaction points (IPs) in the collider, both with and without beam-beam effects.

Parameters:

Name Type Description Default
collider

(xt.Multiline): The collider object configured.

 required
l_n_collisions list[int]

A list containing the number of colliding bunches for each IP.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
def record_final_luminosity(self, collider: xt.Multiline, l_n_collisions: list[int]) -> None:
    """
    Records the final luminosity and pile-up for specified interaction points (IPs)
    in the collider, both with and without beam-beam effects.

    Args:
        collider : (xt.Multiline): The collider object configured.
        l_n_collisions (list[int]): A list containing the number of colliding bunches for each
            IP.

    Returns:
        None
    """
    # Define IPs in which the luminosity will be computed
    l_ip = ["ip1", "ip2", "ip5", "ip8"]

    # Ensure that the final number of particles per bunch is defined, even
    # if the leveling has been done by separation
    if "final_num_particles_per_bunch" not in self.config_beambeam:
        self.config_beambeam["final_num_particles_per_bunch"] = self.config_beambeam[
            "num_particles_per_bunch"
        ]

    def _twiss_and_compute_lumi(collider, l_n_collisions):
        # Loop over each IP and record the luminosity
        twiss_b1 = collider["lhcb1"].twiss()
        twiss_b2 = collider["lhcb2"].twiss()
        l_lumi = []
        l_PU = []
        for n_col, ip in zip(l_n_collisions, l_ip):
            L = xt.lumi.luminosity_from_twiss(  # type: ignore
                n_colliding_bunches=n_col,
                num_particles_per_bunch=self.config_beambeam["final_num_particles_per_bunch"],
                ip_name=ip,
                nemitt_x=self.config_beambeam["nemitt_x"],
                nemitt_y=self.config_beambeam["nemitt_y"],
                sigma_z=self.config_beambeam["sigma_z"],
                twiss_b1=twiss_b1,
                twiss_b2=twiss_b2,
                crab=self.crab,
            )
            PU = compute_PU(
                L,
                n_col,
                twiss_b1["T_rev0"],
                cross_section=self.config_beambeam["cross_section"],
            )

            l_lumi.append(L)
            l_PU.append(PU)

        return l_lumi, l_PU

    # Get the final luminosity in all IPs, without beam-beam
    collider.vars["beambeam_scale"] = 0
    l_lumi, l_PU = _twiss_and_compute_lumi(collider, l_n_collisions)

    # Update configuration
    for ip, L, PU in zip(l_ip, l_lumi, l_PU):
        self.config_beambeam[f"luminosity_{ip}_without_beam_beam"] = float(L)
        self.config_beambeam[f"Pile-up_{ip}_without_beam_beam"] = float(PU)

    # Get the final luminosity in all IPs, with beam-beam
    collider.vars["beambeam_scale"] = 1
    l_lumi, l_PU = _twiss_and_compute_lumi(collider, l_n_collisions)

    # Update configuration
    for ip, L, PU in zip(l_ip, l_lumi, l_PU):
        self.config_beambeam[f"luminosity_{ip}_with_beam_beam"] = float(L)
        self.config_beambeam[f"Pile-up_{ip}_with_beam_beam"] = float(PU)

return_fingerprint(collider, line_name='lhcb1') staticmethod

Generate a detailed fingerprint of the specified collider line. Useful to compare two colliders.

Parameters:

Name Type Description Default
collider Multiline

The collider object containing the line data.

 required
line_name str

The name of the line to analyze within the collider. Default to "lhcb1".

 'lhcb1'

Returns:

Name Type Description
str str

A formatted string containing detailed information about the collider line, including: - Installed element types - Tunes and chromaticity - Synchrotron tune and slip factor - Twiss parameters and phases at interaction points (IPs) - Dispersion and crab dispersion at IPs - Amplitude detuning coefficients - Non-linear chromaticity - Tunes and momentum compaction vs delta
Source code in study_da/generate/master_classes/xsuite_collider.py 
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
@staticmethod
def return_fingerprint(collider, line_name="lhcb1") -> str:
    """
    Generate a detailed fingerprint of the specified collider line. Useful to compare two
    colliders.

    Args:
        collider (xt.Multiline): The collider object containing the line data.
        line_name (str): The name of the line to analyze within the collider. Default to "lhcb1".

    Returns:
        str:
            A formatted string containing detailed information about the collider line, including:
            - Installed element types
            - Tunes and chromaticity
            - Synchrotron tune and slip factor
            - Twiss parameters and phases at interaction points (IPs)
            - Dispersion and crab dispersion at IPs
            - Amplitude detuning coefficients
            - Non-linear chromaticity
            - Tunes and momentum compaction vs delta
    """
    line = collider[line_name]

    tw = line.twiss()
    tt = line.get_table()

    det = line.get_amplitude_detuning_coefficients(a0_sigmas=0.1, a1_sigmas=0.2, a2_sigmas=0.3)

    det_table = xt.Table(
        {
            "name": np.array(list(det.keys())),
            "value": np.array(list(det.values())),
        }
    )

    nl_chrom = line.get_non_linear_chromaticity(
        delta0_range=(-2e-4, 2e-4), num_delta=5, fit_order=3
    )

    out = ""

    out += f"Line: {line_name}\n"
    out += "\n"

    out += "Installed element types:\n"
    out += repr([nn for nn in sorted(list(set(tt.element_type))) if len(nn) > 0]) + "\n"
    out += "\n"

    out += f'Tunes:        Qx  = {tw["qx"]:.5f}       Qy = {tw["qy"]:.5f}\n'
    out += f"""Chromaticity: Q'x = {tw["dqx"]:.2f}     Q'y = """ + f'{tw["dqy"]:.2f}\n'
    out += f'c_minus:      {tw["c_minus"]:.5e}\n'
    out += "\n"

    out += f'Synchrotron tune: {tw["qs"]:5e}\n'
    out += f'Slip factor:      {tw["slip_factor"]:.5e}\n'
    out += "\n"

    out += "Twiss parameters and phases at IPs:\n"
    out += (
        tw.rows["ip.*"]
        .cols["name s betx bety alfx alfy mux muy"]
        .show(output=str, max_col_width=int(1e6), digits=8)
    )
    out += "\n\n"

    out += "Dispersion at IPs:\n"
    out += (
        tw.rows["ip.*"]
        .cols["name s dx dy dpx dpy"]
        .show(output=str, max_col_width=int(1e6), digits=8)
    )
    out += "\n\n"

    out += "Crab dispersion at IPs:\n"
    out += (
        tw.rows["ip.*"]
        .cols["name s dx_zeta dy_zeta dpx_zeta dpy_zeta"]
        .show(output=str, max_col_width=int(1e6), digits=8)
    )
    out += "\n\n"

    out += "Amplitude detuning coefficients:\n"
    out += det_table.show(output=str, max_col_width=int(1e6), digits=6)
    out += "\n\n"

    out += "Non-linear chromaticity:\n"
    out += f'dnqx = {list(nl_chrom["dnqx"])}\n'
    out += f'dnqy = {list(nl_chrom["dnqy"])}\n'
    out += "\n\n"

    out += "Tunes and momentum compaction vs delta:\n"
    out += nl_chrom.show(output=str, max_col_width=int(1e6), digits=6)
    out += "\n\n"

    return out

set_filling_and_bunch_tracked(ask_worst_bunch=False)

Sets the filling scheme and determines the bunch to be tracked for beam-beam interactions.

This method performs the following steps: 1. Retrieves the filling scheme path from the configuration. 2. Checks if the filling scheme path needs to be obtained from the template schemes. 3. Loads and verifies the filling scheme, potentially converting it if necessary. 4. Updates the configuration with the correct filling scheme path. 5. Determines the number of long-range encounters to consider. 6. If the bunch number for beam 1 is not provided, it identifies the bunch with the largest number of long-range interactions. - If ask_worst_bunch is True, prompts the user to confirm or provide a bunch number. - Otherwise, automatically selects the worst bunch. 7. If the bunch number for beam 2 is not provided, it automatically selects the worst bunch.

Parameters:

Name Type Description Default
ask_worst_bunch bool

If True, prompts the user to confirm or provide the bunch number for beam 1. Defaults to False.

 False

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
def set_filling_and_bunch_tracked(self, ask_worst_bunch: bool = False) -> None:
    """
    Sets the filling scheme and determines the bunch to be tracked for beam-beam interactions.

    This method performs the following steps:
    1. Retrieves the filling scheme path from the configuration.
    2. Checks if the filling scheme path needs to be obtained from the template schemes.
    3. Loads and verifies the filling scheme, potentially converting it if necessary.
    4. Updates the configuration with the correct filling scheme path.
    5. Determines the number of long-range encounters to consider.
    6. If the bunch number for beam 1 is not provided, it identifies the bunch with the largest
    number of long-range interactions.
       - If `ask_worst_bunch` is True, prompts the user to confirm or provide a bunch number.
       - Otherwise, automatically selects the worst bunch.
    7. If the bunch number for beam 2 is not provided, it automatically selects the worst bunch.

    Args:
        ask_worst_bunch (bool): If True, prompts the user to confirm or provide the bunch number
            for beam 1. Defaults to False.

    Returns:
        None
    """
    # Get the filling scheme path
    filling_scheme_path = self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"]

    # Check if the filling scheme path must be obtained from the template schemes
    scheme_folder = (
        pathlib.Path(__file__).parent.parent.parent.resolve().joinpath("assets/filling_schemes")
    )
    if filling_scheme_path in os.listdir(scheme_folder):
        filling_scheme_path = str(scheme_folder.joinpath(filling_scheme_path))
        self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"] = filling_scheme_path

    # Load and check filling scheme, potentially convert it
    filling_scheme_path = load_and_check_filling_scheme(filling_scheme_path)

    # Correct filling scheme in config, as it might have been converted
    self.config_beambeam["mask_with_filling_pattern"]["pattern_fname"] = filling_scheme_path

    # Get number of LR to consider
    n_LR = self.config_beambeam["num_long_range_encounters_per_side"]["ip1"]

    # If the bunch number is None, the bunch with the largest number of long-range interactions is used
    if self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"] is None:
        # Case the bunch number has not been provided
        worst_bunch_b1 = get_worst_bunch(
            filling_scheme_path, number_of_LR_to_consider=n_LR, beam="beam_1"
        )
        if ask_worst_bunch:
            while self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"] is None:
                bool_inp = input(
                    "The bunch number for beam 1 has not been provided. Do you want to use the"
                    " bunch with the largest number of long-range interactions? It is the bunch"
                    " number " + str(worst_bunch_b1) + " (y/n): "
                )
                if bool_inp == "y":
                    self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"] = (
                        worst_bunch_b1
                    )
                elif bool_inp == "n":
                    self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"] = int(
                        input("Please enter the bunch number for beam 1: ")
                    )
        else:
            self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b1"] = worst_bunch_b1

    if self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b2"] is None:
        worst_bunch_b2 = get_worst_bunch(
            filling_scheme_path, number_of_LR_to_consider=n_LR, beam="beam_2"
        )
        # For beam 2, just select the worst bunch by default
        self.config_beambeam["mask_with_filling_pattern"]["i_bunch_b2"] = worst_bunch_b2

set_knobs(collider)

Set all knobs for the collider, including crossing angles, dispersion correction, RF, crab cavities, experimental magnets, etc.

Parameters:

Name Type Description Default
collider Multiline

The collider object to which the knob settings will be applied.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
def set_knobs(self, collider: xt.Multiline) -> None:
    """
    Set all knobs for the collider, including crossing angles, dispersion correction,
    RF, crab cavities, experimental magnets, etc.

    Args:
        collider (xt.Multiline): The collider object to which the knob settings will be applied.

    Returns:
        None
    """
    # Set all knobs (crossing angles, dispersion correction, rf, crab cavities,
    # experimental magnets, etc.)
    for kk, vv in self.config_knobs_and_tuning["knob_settings"].items():
        collider.vars[kk] = vv

    # Crab fix (if needed)
    if self.ver_hllhc_optics is not None and self.ver_hllhc_optics == 1.3:
        apply_crab_fix(collider, self.config_knobs_and_tuning)

update_configuration_knob(collider, dictionnary, knob_name) staticmethod

Updates the given dictionary with the final value of a specified knob from the collider.

Parameters:

Name Type Description Default
collider Multiline

The collider object containing various variables.

 required
dictionnary dict

The dictionary to be updated with the knob's final value.

 required
knob_name str

The name of the knob whose value is to be retrieved and stored.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
@staticmethod
def update_configuration_knob(
    collider: xt.Multiline, dictionnary: dict, knob_name: str
) -> None:
    """
    Updates the given dictionary with the final value of a specified knob from the collider.

    Args:
        collider (xt.Multiline): The collider object containing various variables.
        dictionnary (dict): The dictionary to be updated with the knob's final value.
        knob_name (str): The name of the knob whose value is to be retrieved and stored.

    Returns:
        None
    """
    if knob_name in collider.vars.keys():
        dictionnary[f"final_{knob_name}"] = float(collider.vars[knob_name]._value)
    else:
        logging.warning(f"Knob {knob_name} not found in the collider")

write_collider_to_disk(collider, full_configuration)

Writes the collider object to disk in JSON format if the save_output_collider flag is set.

Parameters:

Name Type Description Default
collider Collider

The collider object to be saved.

 required
full_configuration dict

The full configuration dictionary to be deep-copied into the collider's metadata.

 required

Returns:

Type Description
None

None
Source code in study_da/generate/master_classes/xsuite_collider.py 
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
def write_collider_to_disk(self, collider, full_configuration) -> None:
    """
    Writes the collider object to disk in JSON format if the save_output_collider flag is set.

    Args:
        collider (Collider): The collider object to be saved.
        full_configuration (dict): The full configuration dictionary to be deep-copied into the
            collider's metadata.

    Returns:
        None
    """
    if self.save_output_collider:
        logging.info("Saving collider as json")
        if (
            hasattr(collider, "metadata")
            and collider.metadata is not None
            and isinstance(collider.metadata, dict)
        ):
            collider.metadata.update(copy.deepcopy(full_configuration))
        else:
            collider.metadata = copy.deepcopy(full_configuration)
        collider.to_json(self.path_collider_file_for_tracking_as_output)

        # Compress the collider file to zip to ease the load on afs
        if self.compress:
            compress_and_write(self.path_collider_file_for_tracking_as_output)

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()

find_item_in_dic(obj, key)

Find an item in a nested dictionary.

Parameters:

Name Type Description Default
obj dict

The nested dictionary.

 required
key str

The key to find in the nested dictionary.

 required

Returns:

Name Type Description
Any Any

The value corresponding to the key in the nested dictionary.
Source code in study_da/utils/dic_utils.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def find_item_in_dic(obj: dict, key: str) -> Any:
    """Find an item in a nested dictionary.

    Args:
        obj (dict): The nested dictionary.
        key (str): The key to find in the nested dictionary.

    Returns:
        Any: The value corresponding to the key in the nested dictionary.

    """
    if key in obj:
        return obj[key]
    for v in obj.values():
        if isinstance(v, dict):
            item = find_item_in_dic(v, key)
            if item is not None:
                return item

load_dic_from_path(path, ryaml=None)

Load a dictionary from a yaml file.

Parameters:

Name Type Description Default
path str

The path to the yaml file.

 required
ryaml YAML

The yaml reader.

 None

Returns:

Type Description
tuple[dict, YAML]

tuple[dict, ruamel.yaml.YAML]: The dictionary and the yaml reader.
Source code in study_da/utils/dic_utils.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def load_dic_from_path(
    path: str, ryaml: ruamel.yaml.YAML | None = None
) -> tuple[dict, ruamel.yaml.YAML]:
    """Load a dictionary from a yaml file.

    Args:
        path (str): The path to the yaml file.
        ryaml (ruamel.yaml.YAML): The yaml reader.

    Returns:
        tuple[dict, ruamel.yaml.YAML]: The dictionary and the yaml reader.

    """

    if ryaml is None:
        # Initialize yaml reader
        ryaml = ruamel.yaml.YAML()

    # Load dic
    with open(path, "r") as fid:
        dic = ryaml.load(fid)

    return dic, ryaml

nested_get(dic, keys)

Get the value from a nested dictionary using a list of keys.

Parameters:

Name Type Description Default
dic dict

The nested dictionary.

 required
keys list

The list of keys to traverse the nested dictionary.

 required

Returns:

Name Type Description
Any Any

The value corresponding to the keys in the nested dictionary.
Source code in study_da/utils/dic_utils.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
def nested_get(dic: dict, keys: list) -> Any:
    # Adapted from https://stackoverflow.com/questions/14692690/access-nested-dictionary-items-via-a-list-of-keys
    """Get the value from a nested dictionary using a list of keys.

    Args:
        dic (dict): The nested dictionary.
        keys (list): The list of keys to traverse the nested dictionary.

    Returns:
        Any: The value corresponding to the keys in the nested dictionary.

    """
    for key in keys:
        dic = dic[key]
    return dic

nested_set(dic, keys, value)

Set a value in a nested dictionary using a list of keys.

Parameters:

Name Type Description Default
dic dict

The nested dictionary.

 required
keys list

The list of keys to traverse the nested dictionary.

 required
value Any

The value to set in the nested dictionary.

 required

Returns:

Type Description
None

None
Source code in study_da/utils/dic_utils.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def nested_set(dic: dict, keys: list, value: Any) -> None:
    """Set a value in a nested dictionary using a list of keys.

    Args:
        dic (dict): The nested dictionary.
        keys (list): The list of keys to traverse the nested dictionary.
        value (Any): The value to set in the nested dictionary.

    Returns:
        None

    """
    for key in keys[:-1]:
        dic = dic.setdefault(key, {})
    dic[keys[-1]] = value

set_item_in_dic(obj, key, value, found=False)

Set an item in a nested dictionary.

Parameters:

Name Type Description Default
obj dict

The nested dictionary.

 required
key str

The key to set in the nested dictionary.

 required
value Any

The value to set in the nested dictionary.

 required
found bool

Whether the key has been found in the nested dictionary.

 False

Returns:

Type Description
None

None
Source code in study_da/utils/dic_utils.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def set_item_in_dic(obj: dict, key: str, value: Any, found: bool = False) -> None:
    """Set an item in a nested dictionary.

    Args:
        obj (dict): The nested dictionary.
        key (str): The key to set in the nested dictionary.
        value (Any): The value to set in the nested dictionary.
        found (bool): Whether the key has been found in the nested dictionary.

    Returns:
        None

    """
    if key in obj:
        if found:
            raise ValueError(f"Key {key} found more than once in the nested dictionary.")

        obj[key] = value
        found = True
    for v in obj.values():
        if isinstance(v, dict):
            set_item_in_dic(v, key, value, found)

write_dic_to_path(dic, path, ryaml=None)

Write a dictionary to a yaml file.

Parameters:

Name Type Description Default
dic dict

The dictionary to write.

 required
path str

The path to the yaml file.

 required
ryaml YAML

The yaml reader.

 None

Returns:

Type Description
None

None
Source code in study_da/utils/dic_utils.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def write_dic_to_path(dic: dict, path: str, ryaml: ruamel.yaml.YAML | None = None) -> None:
    """Write a dictionary to a yaml file.

    Args:
        dic (dict): The dictionary to write.
        path (str): The path to the yaml file.
        ryaml (ruamel.yaml.YAML): The yaml reader.

    Returns:
        None

    """

    if ryaml is None:
        # Initialize yaml reader
        ryaml = ruamel.yaml.YAML()

    # Write dic
    with open(path, "w") as fid:
        ryaml.dump(dic, fid)
        # Force os to write to disk now, to avoid race conditions
        fid.flush()
        os.fsync(fid.fileno())