Utilities

Helper functions to prepare the twiss dataframes for use with the optics functions as well as reusable utilities, that are needed within multiple optics calculations.

optics_functions.utils.add_missing_columns(df: DataFrame, columns: Iterable) → DataFrame[source]

Check if columns are in df and add them all zero if not.

optics_functions.utils.dphi(data: ndarray, q: float) → ndarray[source]

Return dphi from phase advances in data, see Eq. (8) in [FranchiAnalyticFormulas2017] .

Parameters:

  • data (DataFrame, Series) -- Phase-Advance data.

  • q -- Tune.

Returns:

dphi’s in matrix of shape of data.

optics_functions.utils.dphi_at_element(df: DataFrame, element: str, qx: float, qy: float) → dict[source]

Return dphis for both planes at the given element. See Eq. (8) in [FranchiAnalyticFormulas2017] .

Parameters:

  • df (DataFrame) -- DataFrame containing the Phase-Advance columns for both planes.

  • element -- Element at which dphi is calculated (must be in the index of df).

  • qx (float) -- Tune in X-Plane.

  • qy (float) -- Tune in Y-Plane.

Returns:

dict of planes with DPhi[i] = Phi[element] - Phi[i] (+ Tune if Phi[i] > Phi[element])

optics_functions.utils.get_all_phase_advances(df: DataFrame) → dict[source]

Calculate phase advances between all elements. Will result in a elements x elements matrix, that might be very large!

Parameters:

df (DataFrame) -- DataFrame with phase advance columns (MUX, MUY).

Returns:

Dictionary with DataFrame matrices similar to dmu(i,j) = mu(j) - mu(i).

optics_functions.utils.get_format_keys(format_str: str) → list[source]

Get keys from format string. Unnamed placeholders are returned as empty strings.

optics_functions.utils.i_pow(n: int) → complex[source]

Calculates i**n in a quick and exact way.

optics_functions.utils.merge_complex_columns(df: DataFrame, columns: Sequence[str], drop: bool = True) → TfsDataFrame[source]

Merges the given real and imag columns into complex columns.

Parameters:

  • df (pd.DataFrame) -- DataFrame containing the original columns.

  • columns (Sequence[str]) -- List of complex columns names to be created.

  • drop (bool) -- Original columns are not present in resulting DataFrame.

Returns:

Original TfsDataFrame with added columns.

optics_functions.utils.prepare_twiss_dataframe(df_twiss: TfsDataFrame, df_errors: DataFrame = None, invert_signs_madx: bool = False, max_order: int = 16, join: str = 'inner') → TfsDataFrame[source]

Prepare dataframe to use with the optics functions.

  • Adapt Beam 4 signs.

  • Add missing K(S)L and orbit columns.

  • Merge optics and error dataframes (add values).

Parameters:

  • df_twiss (TfsDataFrame) -- Twiss-optics DataFrame.

  • df_errors (DataFrame) -- Twiss-errors DataFrame (optional).

  • invert_signs_madx (bool) -- Inverts signs after the madx-convention for beam 4. That is, if you use beam 4 you should set this flag to True to convert beam 4 to beam 2 signs.

  • max_order (int) -- Maximum field order to be still included (1==Dipole).

  • join (str) -- How to join elements of optics and errors. “inner” or “outer”.

Returns:

TfsDataFrame with necessary columns added. If a merge happened, only the necessary columns are present.

optics_functions.utils.seq2str(sequence: Iterable) → str[source]

Converts an Iterable to string of its comma separated elements.

optics_functions.utils.set_name_index(df: DataFrame, df_name='') → DataFrame[source]

Sets the NAME column as index if present and checks for string index.

optics_functions.utils.split_complex_columns(df: DataFrame, columns: Sequence[str], drop: bool = True) → TfsDataFrame[source]

Splits the given complex columns into two real-values columns containing the real and imaginary parts of the original columns.

Parameters:

  • df (TfsDataFrame) -- DataFrame containing the original columns.

  • columns (Sequence[str]) -- List of column names to be replaced.

  • drop (bool) -- Original columns are not present in resulting DataFrame.

Returns:

Original TfsDataFrame with added columns.

optics_functions.utils.switch_signs_for_beam4(df_twiss: DataFrame, df_errors: DataFrame = None) → Tuple[TfsDataFrame, TfsDataFrame][source]

Switch the signs for Beam 4 optics. This is due to the switch in direction for this beam and (anti-) symmetry after a rotation of 180deg around the y-axis of magnets, combined with the fact that the KL values in MAD-X twiss do not change sign, but in the errors they do (otherwise it would compensate). Magnet orders that show anti-symmetry are: a1 (K0SL), b2 (K1L), a3 (K2SL), b4 (K3L) etc. Also the sign for (delta) X is switched back to have the same orientation as beam2.

optics_functions.utils.tau(data: ndarray, q: float) → ndarray[source]

Return tau from phase advances in data, see Eq. (16) in [FranchiAnalyticFormulas2017] .

Parameters:

  • data (DataFrame, Series, Array) -- Phase-Advance data.

  • q -- Tune.

Returns:

tau’s in matrix of shape of data.

optics_functions.utils.timeit(text: str = 'Time used {:.3f}s', print_fun=<bound method Logger.debug of <Logger optics_functions.utils (WARNING)>>)[source]

Timing context with logging/printing output.

Parameters:

  • text (str) -- Text to print. If it contains an unnamed format key, this is filled with the time. Otherwise text is assumed to be description of what is happening in this context.

  • print_fun (function) -- Function used to print the final text.