def _get_data(
twiss: tfs.TfsDataFrame,
model: tfs.TfsDataFrame = None,
add_coupling: bool = False
) -> Tuple[tfs.TfsDataFrame, tfs.TfsDataFrame]:
""" Get's the input data as TfsDataFrames. """
LOG.debug("Loading data.")
def try_reading(df_or_path):
try:
return tfs.read(df_or_path, index=NAME)
except TypeError:
return df_or_path
twiss = try_reading(twiss)
if add_coupling:
twiss = add_coupling_to_model(twiss)
if model is None:
model = twiss.copy()
else:
model = try_reading(model)
if add_coupling:
model = add_coupling_to_model(model)
return twiss, model
def prepare_twiss_dataframe(
df_twiss: TfsDataFrame,
df_errors: pd.DataFrame = None,
invert_signs_madx: bool = False,
max_order: int = 16,
join: str = "inner",
) -> TfsDataFrame:
"""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).
Args:
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.
"""
df_twiss = df_twiss.copy() # As data is moved around
if invert_signs_madx:
df_twiss, df_errors = switch_signs_for_beam4(df_twiss, df_errors)
df_twiss = set_name_index(df_twiss, "twiss")
k_columns = [f"K{n}{s}L" for n in range(max_order) for s in ("S", "")]
orbit_columns = list(PLANES)
if df_errors is None:
return add_missing_columns(df_twiss, k_columns + orbit_columns)
# Merge Dataframes
df_errors = df_errors.copy()
df_errors = set_name_index(df_errors, "error")
index = df_twiss.index.join(df_errors.index, how=join)
if not (set(index) - set(df_twiss.index)):
df = df_twiss.loc[index, :]
else:
df = TfsDataFrame(index=index, headers=df_twiss.headers.copy())
# Merge S column and set zeros in dfs for addition, where elements are missing
for df_self, df_other in ((df_twiss, df_errors), (df_errors,
df_twiss)):
df.loc[df_self.index, S] = df_self[S]
for new_indx in df_other.index.difference(df_self.index):
df_self.loc[new_indx, :] = 0
df = df.sort_values(by=S)
df_twiss = add_missing_columns(df_twiss, k_columns + list(PLANES))
df_errors = add_missing_columns(df_errors,
k_columns + [f"{D}{X}", f"{D}{Y}"])
df_errors = df_errors.rename(columns={f"{D}{p}": p for p in PLANES})
add_columns = k_columns + orbit_columns
df.loc[:, add_columns] = df_twiss[add_columns] + df_errors[add_columns]
for name, df_old in (("twiss", df_twiss), ("errors", df_errors)):
dropped_columns = set(df_old.columns) - set(df.columns)
if dropped_columns:
LOG.warning(
f"The following {name}-columns were dropped on merge: {seq2str(dropped_columns)}"
)
return df