def test_write_compression(self, tmp_path: Path):
"""Test that compression works and compressed files are readable."""
n = 1000
df_example = TfsDataFrame(
data=np.zeros([n, n]), # highly compressible data
headers={"Random": "Data"})
out_file = tmp_path / "data_frame.h5"
write_hdf(out_file, df_example, complevel=0)
assert out_file.is_file()
out_file_compressed = tmp_path / "data_frame_comp.h5"
write_hdf(out_file_compressed, df_example, complevel=9)
assert out_file_compressed.is_file()
assert out_file.stat().st_size > out_file_compressed.stat().st_size
df_read = read_hdf(out_file)
assert_tfs_frame_equal(df_example, df_read)
df_read_compressed = read_hdf(out_file_compressed)
assert_tfs_frame_equal(df_example, df_read_compressed)
def _create_stats_df(df: pd.DataFrame,
parameter: str,
global_index: Any = None) -> TfsDataFrame:
""" Calculates the stats over a given parameter.
Note: Could be refactored to use `group_by`.
Args:
df (DataFrame): DataFrame containing the DA information over all seeds.
parameter (str): The parameter over which we want to average, i.e.
SEED or ANGLE.
global_index (Any): identifier to use as a global index, i.e. the statistics
over all entries are stored here. (e.g. '0' for SEEDs)
"""
operation_map = DotDict({
MEAN: np.mean,
STD: np.std,
MIN: np.min,
MAX: np.max
})
pre_index = [] if global_index is None else [global_index]
index = sorted(set(df[parameter]))
n_total = sum(df[parameter] == index[0])
df_stats = TfsDataFrame(
index=pre_index + index,
columns=[
f"{fun}{al}" for al in DA_COLUMNS
for fun in list(operation_map.keys()) + [N]
],
)
df_stats.headers[HEADER_INFO] = INFO.format(over=OVER_WHICH[parameter],
per=parameter.lower(),
n=n_total)
df_stats.headers[HEADER_NTOTAL] = n_total
for col_da in DA_COLUMNS:
for idx in index:
mask = (df[parameter] == idx) & (df[col_da] != 0)
df_stats.loc[idx, f"{N}{col_da}"] = sum(mask)
for name, operation in operation_map.items():
df_stats.loc[idx,
f"{name}{col_da}"] = operation(df.loc[mask,
col_da])
for name, operation in operation_map.get_subdict([MIN,
MAX]).items():
df_stats.loc[idx, f"{name}{AMP}"] = operation(
df.loc[mask, f"{name}{AMP}"])
if global_index is not None:
# Note: could be done over df_stats for MEAN, MIN and MAX, but not STD
mask = df[col_da] != 0
df_stats.loc[global_index, f"{N}{col_da}"] = sum(mask)
# Global MEAN, MIN, MAX Dynamic Aperture
for name, operation in operation_map.get_subdict(
[MEAN, MIN, MAX, STD]).items():
df_stats.loc[global_index,
f"{name}{col_da}"] = operation(df.loc[mask,
col_da])
# Global MIN, MAX Amplitudes
for name, operation in operation_map.get_subdict([MIN,
MAX]).items():
df_stats.loc[global_index, f"{name}{AMP}"] = operation(
df.loc[mask, f"{name}{AMP}"] # min(MINA) and max(MAXA)
)
df_stats.headers[HEADER_HINT] = HINT.format(param=parameter,
val=global_index)
return df_stats
def calculate_rdts(df: TfsDataFrame,
rdts: Sequence[str],
qx: float = None,
qy: float = None,
feeddown: int = 0,
complex_columns: bool = True,
loop_phases: bool = False,
hamiltionian_terms: bool = False) -> TfsDataFrame:
""" Calculates the Resonance Driving Terms.
Eq. (A8) in [FranchiAnalyticFormulas2017]_ .
Args:
df (TfsDataFrame): Twiss Dataframe.
rdts (Sequence): List of rdt-names to calculate.
qx (float): Tune in X-Plane (if not given, header df.Q1 is assumed present).
qy (float): Tune in Y-Plane (if not given, header df.Q2 is assumed present).
feeddown (int): Levels of feed-down to include.
complex_columns (bool): Output complex values in single column of type complex.
If ``False``, split complex columns into two real-valued columns.
loop_phases (bool): Loop over elements when calculating phase-advances.
Might be slower for small number of elements, but
allows for large (e.g. sliced) optics.
hamiltionian_terms (bool): Add the hamiltonian terms to the result dataframe.
Returns:
New TfsDataFrame with RDT columns.
"""
LOG.info(f"Calculating RDTs: {seq2str(rdts):s}.")
with timeit("RDT calculation", print_fun=LOG.debug):
df_res = TfsDataFrame(index=df.index)
if qx is None:
qx = df.headers[f"{TUNE}1"]
if qy is None:
qy = df.headers[f"{TUNE}2"]
if not loop_phases:
phase_advances = get_all_phase_advances(df) # might be huge!
for rdt in rdts:
rdt = rdt.upper()
if len(rdt) != 5 or rdt[0] != 'F':
raise ValueError(
f"'{rdt:s}' does not seem to be a valid RDT name.")
j, k, l, m = [int(i) for i in rdt[1:]]
conj_rdt = jklm2str(k, j, m, l)
if conj_rdt in df_res:
df_res[rdt] = np.conjugate(df_res[conj_rdt])
else:
with timeit(f"calculating {rdt}", print_fun=LOG.debug):
n = j + k + l + m
jk, lm = j + k, l + m
if n <= 1:
raise ValueError(
f"The RDT-order has to be >1 but was {n:d} for {rdt:s}"
)
denom_h = 1. / (factorial(j) * factorial(k) *
factorial(l) * factorial(m) * (2**n))
denom_f = 1. / (1. - np.exp(PI2I * ((j - k) * qx +
(l - m) * qy)))
betax = df[f"{BETA}{X}"]**(jk / 2.)
betay = df[f"{BETA}{Y}"]**(lm / 2.)
# Magnetic Field Strengths with Feed-Down
dx_idy = df[X] + 1j * df[Y]
k_complex = pd.Series(
0j, index=df.index
) # Complex sum of strenghts (from K_n + iJ_n) and feeddown to them
for q in range(feeddown + 1):
n_mad = n + q - 1
kl_iksl = df[f"K{n_mad:d}L"] + 1j * df[f"K{n_mad:d}SL"]
k_complex += (kl_iksl * (dx_idy**q)) / factorial(q)
# real(i**lm * k+ij) is equivalent to Omega-function in paper, see Eq.(A11)
# pd.Series is needed here, as np.real() returns numpy-array
k_real = pd.Series(np.real(i_pow(lm) * k_complex),
index=df.index)
sources = df.index[
k_real !=
0] # other elements do not contribute to integral, speedup summations
if not len(sources):
LOG.warning(
f"No sources found for {rdt}. RDT will be zero.")
df_res[rdt] = 0j
if hamiltionian_terms:
df_res[f2h(rdt)] = 0j
continue
# calculate hamiltionian-numerator part
h_terms = -k_real.loc[sources] * betax.loc[
sources] * betay.loc[sources]
if loop_phases:
# do loop over elements to not have elements x elements Matrix in memory
h_jklm = pd.Series(index=df.index, dtype=complex)
for element in df.index:
# calculate dphi from all sources to the element
# index-intersection keeps `element` at correct place in index
sources_plus = df.index.intersection(
sources.union([element]))
dphis = dphi_at_element(df.loc[sources_plus, :],
element, qx, qy)
phase_term = np.exp(
PI2I * ((j - k) * dphis[X].loc[sources] +
(l - m) * dphis[Y].loc[sources]))
h_jklm[element] = (h_terms *
phase_term).sum() * denom_h
else:
phx = dphi(phase_advances['X'].loc[sources, :], qx)
phy = dphi(phase_advances['Y'].loc[sources, :], qy)
phase_term = np.exp(PI2I * ((j - k) * phx +
(l - m) * phy))
h_jklm = phase_term.multiply(
h_terms, axis="index").sum(
axis="index").transpose() * denom_h
df_res[rdt] = h_jklm * denom_f
LOG.info(
f"Average RDT amplitude |{rdt:s}|: {df_res[rdt].abs().mean():g}"
)
if hamiltionian_terms:
df_res[f2h(rdt)] = h_jklm
if not complex_columns:
terms = list(rdts)
if hamiltionian_terms:
terms += [f2h(rdt) for rdt in rdts] # F#### -> H####
df_res = split_complex_columns(df_res, terms)
return df_res
def closest_tune_approach(
df: TfsDataFrame, qx: float = None, qy: float = None, method: str = "teapot"
) -> TfsDataFrame:
"""Calculates the closest tune approach from coupling resonances.
A complex F1001 column is assumed to be present in the DataFrame.
This can be calculated by :func:`~optics_functions.rdt.rdts`
:func:`~optics_functions.coupling.coupling_from_rdts` or
:func:`~optics_functions.coupling.coupling_from_cmatrix`.
If F1010 is also present it is used, otherwise it is assumed 0.
The closest tune approach is calculated by means of Eq. (27) in
[CalagaBetatronCoupling2005]_ (method="teapot" or "calaga") by default,
or approximated by
Eq. (1) in [PerssonImprovedControlCoupling2014]_ (method="franchi"),
Eq. (27) in [CalagaBetatronCoupling2005]_ with the Franchi appoximation (method="teapot_franchi"),
Eq. (2) in [PerssonImprovedControlCoupling2014]_ (method="persson"),
the latter without the exp(i(Qx-Qy)s/R) term (method="persson_alt"),
Eq. (14) in [HoydalsvikEvaluationOfTheClosestTuneApproach2021]_ (method="hoydalsvik"),
or the latter without the exp(i(Qx-Qy)s/R) term (method="hoydalsvik_alt").
For "persson[_alt]" and "hoydalsvik[_alt]" methods, also MUX and MUY columns
are needed in the DataFrame as well as LENGTH (of the machine) and S column
for the "persson" and "hoydalsvik" methods.
Args:
df (TfsDataFrame): Twiss Dataframe, needs to have complex-valued F1001 column.
qx (float): Tune in X-Plane (if not given, header df.Q1 is assumed present).
qy (float): Tune in Y-Plane (if not given, header df.Q2 is assumed present).
method (str): Which method to use for evaluation.
Choices: "calaga", "teapot", "franchi", "teapot_franchi",
"persson", "persson_alt", "hoydalsvik" or "hoydalsvik_alt".
Returns:
A new ``TfsDataFrame`` with a closest tune approach (DELTAQMIN) column.
The value is real for "calaga", "teapot", "teapot_franchi" and "franchi"
methods. The actual closest tune approach value is the absolute value
of the mean of this column.
"""
if F1001 not in df.columns:
raise KeyError(f"'{F1001}' column not in dataframe. Needed to calculated closest tune approach.")
method_map = {
"teapot": _cta_teapot, # as named in [HoydalsvikEvaluationOfTheClosestTuneApproach2021]_
"calaga": _cta_teapot, # for compatibility reasons
"teapot_franchi": _cta_teapot_franchi,
"franchi": _cta_franchi,
"persson": _cta_persson,
"persson_alt": _cta_persson_alt,
"hoydalsvik": _cta_hoydalsvik,
"hoydalsvik_alt": _cta_hoydalsvik_alt,
}
if qx is None:
qx = df.headers[f"{TUNE}1"]
if qy is None:
qy = df.headers[f"{TUNE}2"]
qx_frac, qy_frac = qx % 1, qy % 1
check_resonance_relation(df)
dqmin_str = f"{DELTA}{TUNE}{MINIMUM}"
df_res = TfsDataFrame(index=df.index, columns=[dqmin_str])
df_res[dqmin_str] = method_map[method.lower()](df, qx_frac, qy_frac)
LOG.info(f"({method.lower()}) |C-| = {np.abs(df_res[dqmin_str].dropna().mean())}")
return df_res
def test_write_read_spaces_in_strings(_test_file: str):
df = TfsDataFrame(data=["This is", "a test", 'with spaces'], columns=["A"])
write_tfs(_test_file, df)
new = read_tfs(_test_file)
assert_frame_equal(df, new)
def test_fail_on_spaces_headers():
df = TfsDataFrame(headers={"allowed": 1, "not allowed": 2})
with pytest.raises(TfsFormatError):
write_tfs('', df)
def test_fail_on_spaces_columns():
df = TfsDataFrame(columns=["allowed", "not allowed"])
with pytest.raises(TfsFormatError):
write_tfs('', df)
def test_fail_on_non_unique_index():
df = TfsDataFrame(index=["A", "B", "A"])
with pytest.raises(TfsFormatError):
write_tfs('', df)
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
