Skip to content

scheme_utils

This class is used to inspect and compute some properties of the filling scheme.

get_worst_bunch(filling_scheme_path, number_of_LR_to_consider=26, beam='beam_1')

Adapted from https://github.com/PyCOMPLETE/FillingPatterns/blob/5f28d1a99e9a2ef7cc5c171d0cab6679946309e8/fillingpatterns/bbFunctions.py#L233

Given a filling scheme, containing two arrays of booleans representing the trains of bunches for the two beams, this function returns the worst bunch for each beam, according to their collision schedule.

Parameters:

Name Type Description Default
filling_scheme_path str

Path to the filling scheme file.

required
number_of_LR_to_consider int

Number of long range collisions to consider. Defaults to 26.

26
beam str

Beam for which to compute the worst bunch. Defaults to "beam_1".

'beam_1'

Returns:

Name Type Description
int int

The worst bunch for the specified beam.

Source code in study_da/generate/master_classes/scheme_utils.py
def get_worst_bunch(
    filling_scheme_path: str, number_of_LR_to_consider: int = 26, beam="beam_1"
) -> int:
    """
    # Adapted from https://github.com/PyCOMPLETE/FillingPatterns/blob/5f28d1a99e9a2ef7cc5c171d0cab6679946309e8/fillingpatterns/bbFunctions.py#L233
    Given a filling scheme, containing two arrays of booleans representing the trains of bunches for
    the two beams, this function returns the worst bunch for each beam, according to their collision
    schedule.

    Args:
        filling_scheme_path (str): Path to the filling scheme file.
        number_of_LR_to_consider (int): Number of long range collisions to consider. Defaults to 26.
        beam (str): Beam for which to compute the worst bunch. Defaults to "beam_1".

    Returns:
        int: The worst bunch for the specified beam.

    """

    if not filling_scheme_path.endswith(".json"):
        raise ValueError("Only json filling schemes are supported")

    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"])

    # Get bunches index
    B1_bunches_index = np.flatnonzero(array_b1)
    B2_bunches_index = np.flatnonzero(array_b2)

    # Compute the number of long range collisions per bunch
    l_long_range_per_bunch = _compute_LR_per_bunch(
        array_b1, array_b2, B1_bunches_index, B2_bunches_index, number_of_LR_to_consider, beam=beam
    )

    # Get the worst bunch for both beams
    if beam == "beam_1":
        worst_bunch = B1_bunches_index[np.argmax(l_long_range_per_bunch)]
    elif beam == "beam_2":
        worst_bunch = B2_bunches_index[np.argmax(l_long_range_per_bunch)]
    else:
        raise ValueError("beam must be either 'beam_1' or 'beam_2")

    # Need to explicitly convert to int for json serialization
    return int(worst_bunch)

load_and_check_filling_scheme(filling_scheme_path)

Load and check the filling scheme from a JSON file. Convert the filling scheme to the correct format if needed.

Parameters:

Name Type Description Default
filling_scheme_path str

Path to the filling scheme file.

required

Returns:

Name Type Description
str str

Path to the converted filling scheme file.

Source code in study_da/generate/master_classes/scheme_utils.py
def load_and_check_filling_scheme(filling_scheme_path: str) -> str:
    """Load and check the filling scheme from a JSON file. Convert the filling scheme to the correct
    format if needed.

    Args:
        filling_scheme_path (str): Path to the filling scheme file.

    Returns:
        str: Path to the converted filling scheme file.
    """
    if not filling_scheme_path.endswith(".json"):
        raise ValueError("Filling scheme must be in json format")

    # Check that the converted filling scheme doesn't already exist
    filling_scheme_path_converted = filling_scheme_path.replace(".json", "_converted.json")
    if os.path.exists(filling_scheme_path_converted):
        return filling_scheme_path_converted

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

    if "beam1" in d_filling_scheme.keys() and "beam2" in d_filling_scheme.keys():
        # If the filling scheme not already in the correct format, convert
        if "schemebeam1" in d_filling_scheme.keys() or "schemebeam2" in d_filling_scheme.keys():
            d_filling_scheme["beam1"] = d_filling_scheme["schemebeam1"]
            d_filling_scheme["beam2"] = d_filling_scheme["schemebeam2"]
            # Delete all the other keys
            d_filling_scheme = {
                k: v for k, v in d_filling_scheme.items() if k in ["beam1", "beam2"]
            }
            # Dump the dictionary back to the file
            with open(filling_scheme_path_converted, "w") as fid:
                json.dump(d_filling_scheme, fid)

            # Else, do nothing

    else:
        # One can potentially use b1_array, b2_array to scan the bunches later
        b1_array, b2_array = reformat_filling_scheme_from_lpc(
            filling_scheme_path, filling_scheme_path_converted
        )
        filling_scheme_path = filling_scheme_path_converted

    return filling_scheme_path

reformat_filling_scheme_from_lpc(filling_scheme_path, filling_scheme_path_converted)

This function is used to convert the filling scheme from the LPC to the format used in the xtrack library. The filling scheme from the LPC is a list of bunches for each beam, where each bunch is represented by a 1 in the list. The function converts this list to a list of indices of the filled bunches. The function also returns the indices of the filled bunches for each beam.

Parameters:

Name Type Description Default
filling_scheme_path str

Path to the filling scheme file.

required
filling_scheme_path_converted str

Path to the converted filling scheme file.

required

Returns:

Type Description
tuple[ndarray, ndarray]

tuple[np.ndarray, np.ndarray]: Indices of the filled bunches for each beam.

Source code in study_da/generate/master_classes/scheme_utils.py
def reformat_filling_scheme_from_lpc(
    filling_scheme_path: str, filling_scheme_path_converted: str
) -> tuple[np.ndarray, np.ndarray]:
    """
    This function is used to convert the filling scheme from the LPC to the format used in the
    xtrack library. The filling scheme from the LPC is a list of bunches for each beam, where each
    bunch is represented by a 1 in the list. The function converts this list to a list of indices
    of the filled bunches. The function also returns the indices of the filled bunches for each beam.

    Args:
        filling_scheme_path (str): Path to the filling scheme file.
        filling_scheme_path_converted (str): Path to the converted filling scheme file.

    Returns:
        tuple[np.ndarray, np.ndarray]: Indices of the filled bunches for each beam.
    """

    # Load the filling scheme directly if json
    with open(filling_scheme_path, "r") as fid:
        data = json.load(fid)

    # Take the first fill number
    fill_number = list(data["fills"].keys())[0]

    # Do the conversion (Matteo's code)
    B1 = np.zeros(3564)
    B2 = np.zeros(3564)
    l_lines = data["fills"][f"{fill_number}"]["csv"].split("\n")
    for idx, line in enumerate(l_lines):
        # First time one encounters a line with 'Slot' in it, start indexing
        if "Slot" in line:
            # B1 is initially empty
            if np.sum(B1) == 0:
                for line_2 in l_lines[idx + 1 :]:
                    l_line = line_2.split(",")
                    if len(l_line) > 1:
                        slot = l_line[1]
                        B1[int(slot)] = 1
                    else:
                        break

            elif np.sum(B2) == 0:
                for line_2 in l_lines[idx + 1 :]:
                    l_line = line_2.split(",")
                    if len(l_line) > 1:
                        slot = l_line[1]
                        B2[int(slot)] = 1
                    else:
                        break
            else:
                break

    data_json = {"beam1": [int(ii) for ii in B1], "beam2": [int(ii) for ii in B2]}

    with open(filling_scheme_path_converted, "w") as file_bool:
        json.dump(data_json, file_bool)
    return B1, B2