def test_tfs_read_write_read_pathlib_input(_tfs_file_pathlib: pathlib.Path, _test_file: str): original = read_tfs(_tfs_file_pathlib) write_tfs(_test_file, original) new = read_tfs(_test_file) assert_frame_equal(original, new) assert_dict_equal(original.headers, new.headers, compare_keys=True)
def test_tfs_read_write_read(_tfs_file, _test_file): original = read_tfs(_tfs_file) write_tfs(_test_file, original) new = read_tfs(_test_file) assert original.headers == new.headers assert all(original.columns == new.columns) for column in original: assert all(original.loc[:, column] == new.loc[:, column])
def test_tfs_read_write_read_pathlib_input(self, _tfs_file_pathlib: pathlib.Path, tmp_path): original = read_tfs(_tfs_file_pathlib) write_location = tmp_path / "test_file.tfs" write_tfs(write_location, original) new = read_tfs(write_location) assert_frame_equal(original, new) assert_dict_equal(original.headers, new.headers, compare_keys=True)
def merge_tfs(directories: List[pathlib.Path], filename: str) -> tfs.TfsDataFrame: """ Merge different kmod analysis results from a list of directories into a single `TfsDataFrame`. Args: directories (List[pathlib.Path]): list of PosixPath objects to directories holding TFS files with the results of kmod analysis. filename (str): name of the TFS files to look for in the provided directories Returns: A `TfsDataFrame` combining all the loaded files from the provided directories. """ # Combine the data into one tfs new_tfs = tfs.TfsDataFrame() headers = {} for d in sorted(directories): loaded_tfs = tfs.read_tfs(d / filename) headers.update(loaded_tfs.headers) # old headers are lost in `append` new_tfs = new_tfs.append(loaded_tfs, ignore_index=True) new_tfs.headers = headers new_tfs = new_tfs.set_index(NAME) # drop BPMWK and check tfs new_tfs = new_tfs.loc[~new_tfs.index.str.startswith(BPMWK), :] if not new_tfs.index.is_unique: raise KeyError( "Found duplicated entries " f"{', '.join(set(new_tfs.index[new_tfs.index.duplicated()]))}'.") return new_tfs
def test_absent_attributes_and_keys(self, _tfs_file_str: str): test_file = read_tfs(_tfs_file_str, index="NAME") with pytest.raises(AttributeError): _ = test_file.Not_HERE with pytest.raises(KeyError): _ = test_file["Not_HERE"]
def test_warn_unphysical_values(self, caplog): nan_tfs_path = pathlib.Path( __file__).parent / "inputs" / "has_nans.tfs" _ = read_tfs(nan_tfs_path, index="NAME") for record in caplog.records: assert record.levelname == "WARNING" assert "contains non-physical values at Index:" in caplog.text
def test_write_read_spaces_in_strings(self, tmp_path): df = TfsDataFrame(data=["This is", "a test", "with spaces"], columns=["A"]) write_location = tmp_path / "test.tfs" write_tfs(write_location, df) new = read_tfs(write_location) assert_frame_equal(df, new)
def test_merge_kmod_results_commandline(tmp_path): paths = [str(INPUT_DIR / "kmod_ip1"), str(INPUT_DIR / "kmod_ip5")] with cli_args("--kmod_dirs", *paths, "--outputdir", str(tmp_path)): merge_kmod_results.merge_kmod_results() filename = f"{merge_kmod_results.LSA_RESULTS}{merge_kmod_results.EXT}" res_lsa_tfs = tfs.read_tfs(tmp_path / filename, index=merge_kmod_results.NAME) control_tfs = tfs.read_tfs(INPUT_DIR / "lsa_results_merged.tfs", index=merge_kmod_results.NAME) assert_frame_equal(res_lsa_tfs, control_tfs) assert_dict_equal(res_lsa_tfs.headers, control_tfs.headers, compare_keys=True)
def test_tfs_read_str_input(self, _tfs_file_str: str): test_file = read_tfs(_tfs_file_str, index="NAME") assert len(test_file.headers) > 0 assert len(test_file.columns) > 0 assert len(test_file.index) > 0 assert len(str(test_file)) > 0 assert isinstance(test_file.index[0], str)
def test_tfs_write_read(_dataframe: TfsDataFrame, _test_file: str): write_tfs(_test_file, _dataframe) assert pathlib.Path(_test_file).is_file() new = read_tfs(_test_file) assert_frame_equal(_dataframe, new, check_exact=False) # float precision can be an issue assert_dict_equal(_dataframe.headers, new.headers, compare_keys=True)
def test_tfs_write_read_autoindex(_dataframe: TfsDataFrame, _test_file: str): df = _dataframe.set_index("a") df1 = _dataframe.set_index("a") write_tfs(_test_file, df, save_index=True) assert_frame_equal(df, df1) df_read = read_tfs(_test_file) assert_index_equal(df.index, df_read.index, check_exact=False) assert_dict_equal(_dataframe.headers, df_read.headers, compare_keys=True)
def test_tfs_write_read_autoindex(_dataframe, _test_file): df = _dataframe.set_index("a") df1 = _dataframe.set_index("a") write_tfs(_test_file, df, save_index=True) compare_dataframes(df, df1) # writing should not change things df_read = read_tfs(_test_file) assert df_read.index.name == df.index.name assert all((df_read.index.values - df.index.values) <= 1E-12)
def test_write_int_float_columns(self, tmp_path): """This test is here because of numeric conversion bug upon writing back in v2.0.1""" df = TfsDataFrame(data=[[1, 1.0], [2, 2.0], [3, 3.0]], columns=["Int", "Float"]) write_location = tmp_path / "test.tfs" write_tfs(write_location, df) new = read_tfs(write_location) assert_frame_equal(df, new)
def get_ip_positions(path): """ Returns a dict of IP positions from tfs-file of path. Args: path (str): Path to the tfs-file containing IP-positions """ df = tfs.read_tfs(path).set_index('NAME') ip_names = [f"IP{i:d}" for i in range(1, 9)] ip_pos = df.loc[ip_names, 'S'].values return dict(zip(ip_names, ip_pos))
def test_merge_kmod_results(tmp_path): paths = [INPUT_DIR / "kmod_ip1", INPUT_DIR / "kmod_ip5"] res_tfs_passed = merge_kmod_results.merge_kmod_results(kmod_dirs=paths, outputdir=tmp_path) filename = f"{merge_kmod_results.LSA_RESULTS}{merge_kmod_results.EXT}" res_lsa_tfs = tfs.read_tfs(tmp_path / filename, index=merge_kmod_results.NAME) control_tfs = tfs.read_tfs(INPUT_DIR / "lsa_results_merged.tfs", index=merge_kmod_results.NAME) assert_frame_equal(res_lsa_tfs, control_tfs) assert_dict_equal(res_lsa_tfs.headers, control_tfs.headers, compare_keys=True) assert_frame_equal(res_tfs_passed, control_tfs, check_exact=False) assert_dict_equal(res_tfs_passed.headers, control_tfs.headers, compare_keys=True)
def test_write_int_float_str_columns(self, tmp_path): """This test is more of an extension of the test below (this dataframe was not affected by the bug)""" df = TfsDataFrame( data=[[1, 1.0, "one"], [2, 2.0, "two"], [3, 3.0, "three"]], columns=["Int", "Float", "String"], ) write_location = tmp_path / "test.tfs" write_tfs(write_location, df) new = read_tfs(write_location) assert_frame_equal(df, new)
def test_tfs_write_read_pandasdf(_pddataframe: DataFrame, _test_file: str): write_tfs(_test_file, _pddataframe) assert pathlib.Path(_test_file).is_file() new = read_tfs(_test_file) assert_frame_equal( _pddataframe, new, check_exact=False, # float precision can be an issue check_frame_type=False, # read df is TfsDF )
def test_tfs_write_read_autoindex(self, _tfs_dataframe, tmp_path): df = _tfs_dataframe.set_index("a") df1 = _tfs_dataframe.set_index("a") write_location = tmp_path / "test.tfs" write_tfs(write_location, df, save_index=True) assert_frame_equal(df, df1) df_read = read_tfs(write_location) assert_index_equal(df.index, df_read.index, check_exact=False) assert_dict_equal(_tfs_dataframe.headers, df_read.headers, compare_keys=True)
def test_tfs_write_read_pandasdf(self, _pd_dataframe, tmp_path): write_location = tmp_path / "test.tfs" write_tfs(write_location, _pd_dataframe) assert write_location.is_file() new = read_tfs(write_location) assert_frame_equal( _pd_dataframe, new, check_exact=False, # float precision can be an issue check_frame_type=False, # read df is TfsDF )
def test_tfs_write_read(self, _tfs_dataframe, tmp_path): write_location = tmp_path / "test.tfs" write_tfs(write_location, _tfs_dataframe) assert write_location.is_file() new = read_tfs(write_location) assert_frame_equal( _tfs_dataframe, new, check_exact=False) # float precision can be an issue assert_dict_equal(_tfs_dataframe.headers, new.headers, compare_keys=True)
def test_tfs_write_read_no_headers(self, _dataframe_empty_headers: TfsDataFrame, tmp_path): write_location = tmp_path / "test.tfs" write_tfs(write_location, _dataframe_empty_headers) assert write_location.is_file() new = read_tfs(write_location) assert_frame_equal(_dataframe_empty_headers, new, check_exact=False) # float precision assert_dict_equal(_dataframe_empty_headers.headers, new.headers, compare_keys=True)
def test_tfs_read_str_input(_tfs_file_str: str): test_file = read_tfs(_tfs_file_str, index="NAME") assert len(test_file.headers) > 0 assert len(test_file.columns) > 0 assert len(test_file.index) > 0 assert len(str(test_file)) > 0 assert isinstance(test_file.index[0], str) with pytest.raises(AttributeError): test_var = test_file.Not_HERE with pytest.raises(KeyError): test_var = test_file["Not_HERE"]
def test_tfs_read_pathlib_input(self, _tfs_file_pathlib: pathlib.Path): test_file = read_tfs(_tfs_file_pathlib, index="NAME") assert len(test_file.headers) > 0 assert len(test_file.columns) > 0 assert len(test_file.index) > 0 assert len(str(test_file)) > 0 assert isinstance(test_file.index[0], str) with pytest.raises(AttributeError): _ = test_file.Not_HERE with pytest.raises(KeyError): _ = test_file["Not_HERE"]
def clean_columns(files: Sequence[Union[Path, str]], columns: Sequence[str], limit: float = 0.0, backup: bool = True): """ Clean the columns in the given files.""" for file in files: file = Path(file) LOG.info(f"Cleaning {file.name}.") df = tfs.read_tfs(file, index=COL_NAME) for column in columns: df = _filter_by_column(df, column, limit) if backup: _backup_file(file) tfs.write_tfs(file, df, save_index=COL_NAME)
def test_tfs_write_empty_index_dataframe(_test_file: str): df = TfsDataFrame( index=[], columns=["a", "b", "c"], data=numpy.random.rand(0, 3), headers={ "Title": "Tfs Title", "Value": 3.3663 }, ) write_tfs(_test_file, df) assert pathlib.Path(_test_file).is_file() new = read_tfs(_test_file) assert_frame_equal(df, new) assert_dict_equal(df.headers, new.headers, compare_keys=True)
def load_tfs_table(filename): """ Takes a .tfs file and returns the Twiss and SUMM tables. Args: filename(str): Path to .tfs file from a Twiss call. Returns: A Twiss and SUMM table for the given machine and beam. """ # Returns a special DataFrame with the SUMM table stored as individual, # capitalized attributes of the returned DataFrame. special_df = tfs.read_tfs(filename, index='NAME') # Coerce into stadard dictionary and DataFrame d = dict(special_df.headers) df = pd.DataFrame(special_df) return df, d
def test_tfs_write_empty_index_dataframe(self, tmp_path): df = TfsDataFrame( index=[], columns=["a", "b", "c"], data=numpy.random.rand(0, 3), headers={ "Title": "Tfs Title", "Value": 3.3663 }, ) write_location = tmp_path / "test.tfs" write_tfs(write_location, df) assert write_location.is_file() new = read_tfs(write_location) assert_frame_equal(df, new) assert_dict_equal(df.headers, new.headers, compare_keys=True)
def load_tfs(): """ Plots a tfs-file. TODO: * check all tfs files for common columns -> make user choose which column * changemarkers tickbox """ LOG.debug("Load Tfs clicked.") paths = QFileDialog().getOpenFileNames(None, 'Load file(s)', '')[0] fig = None if paths: LOG.info("Files chosen: {:s}".format(", ".join(paths))) if len(paths) > 1: # load all files and check for common columns df_list, common_cols = _get_all_tfs_and_common_columns(paths) column_selector = ColumnSelectorDialog(common_cols, single_line=True) elif len(paths) == 1: # load only one tfs LOG.debug("Loading only one file") try: df = tfs.read_tfs(paths[0]) except TfsFormatError: LOG.error("File '{}' is not of TFS format!".format(paths[0])) else: column_selector = ColumnSelectorDialog(df.columns.tolist()) selected = column_selector.get_selected_columns() if selected: NotImplemented("Plotting not yet implemented") # fig = plot_tfs.plot_single_file( # files=paths, # x_cols=[s["x"] for s in selected], # y_cols=[s["y"] for s in selected], # e_cols=[s["e"] for s in selected], # labels=[s["l"] for s in selected], # no_show=True, # ) return fig LOG.debug("No files chosen.") return None
def _get_all_tfs_and_common_columns(paths): tfs_list = [tfs.read_tfs(p) for p in paths] cols = tfs_list[0].columns for t in tfs_list[1:]: cols = cols.intersection(t.columns) return tfs_list, cols
def test_fail_read_no_colnames(self, _no_colnames_tfs_path, caplog): with pytest.raises(TfsFormatError) as e: _ = read_tfs(_no_colnames_tfs_path) assert "column names" in str(e)