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