def test_warning_on_non_unique_index(self, tmp_path, caplog):
df = TfsDataFrame(index=["A", "B", "A"])
write_tfs(tmp_path / "temporary.tfs", df)
for record in caplog.records:
assert record.levelname == "WARNING"
assert "Non-unique indices found" 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_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_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_raising_on_non_unique_both(self, caplog):
df = TfsDataFrame(index=["A", "B", "A"], columns=["A", "B", "A"])
with pytest.raises(TfsFormatError):
write_tfs("", df, non_unique_behavior="raise")
for record in caplog.records:
assert record.levelname == "WARNING"
assert "Non-unique indices found" in caplog.text # first checked and raised
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_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 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_fail_on_spaces_headers(self, caplog):
caplog.set_level(logging.DEBUG)
df = TfsDataFrame(headers={"allowed": 1, "not allowed": 2})
with pytest.raises(TfsFormatError):
write_tfs("", df)
for record in caplog.records:
assert record.levelname == "DEBUG"
assert "Space(s) found in TFS header names" in caplog.text
def test_fail_on_wrong_column_type(self, caplog):
caplog.set_level(logging.DEBUG)
df = TfsDataFrame(columns=range(5))
with pytest.raises(TfsFormatError):
write_tfs("", df)
for record in caplog.records:
assert record.levelname == "DEBUG"
assert "not of string-type" in caplog.text
def test_fail_on_spaces_columns(self, caplog):
caplog.set_level(logging.DEBUG)
df = TfsDataFrame(columns=["allowed", "not allowed"])
with pytest.raises(TfsFormatError):
write_tfs("", df)
for record in caplog.records:
assert record.levelname == "DEBUG"
assert "Space(s) found in TFS columns" in caplog.text
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 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_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_list_column_dataframe_fails_writes(
self, _list_column_in_dataframe: TfsDataFrame, tmp_path):
list_col_tfs = _list_column_in_dataframe
with pytest.raises(
ValueError
): # truth value of nested can't be assesed in _validate
write_tfs("", list_col_tfs)
del list_col_tfs["d"] # should work without the column of lists
write_location = tmp_path / "test.tfs"
write_tfs(write_location, list_col_tfs)
assert write_location.is_file()
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_dict_column_dataframe_fails_writes(
self, _dict_column_in_dataframe: TfsDataFrame, tmp_path):
dict_col_tfs = _dict_column_in_dataframe
with pytest.raises(
TypeError
): # tries to format dict.__dict__, can't get a % formatter
write_tfs("", dict_col_tfs)
del dict_col_tfs["d"] # should work without the column of dicts
write_location = tmp_path / "test.tfs"
write_tfs(write_location, dict_col_tfs)
assert write_location.is_file()
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 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 _copy_and_modify_linfiles(out_path: Path,
columns=None,
index=None,
by=0.0):
paths = {}
for plane in PLANES:
lin_file_src = _get_inputs_linfile(plane)
lin_file_dst = out_path / lin_file_src.name
paths[plane] = lin_file_dst
if index is not None and columns is not None:
plane_columns = [f"{col}{plane}" for col in columns]
df = tfs.read(lin_file_src)
df.loc[index, plane_columns] = df.loc[index, plane_columns] + by
tfs.write_tfs(lin_file_dst, df)
else:
shutil.copy(lin_file_src, lin_file_dst)
return paths
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 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 test_madx_reads_written_tfsdataframes(self, _bigger_tfs_dataframe,
tmp_path):
dframe = _bigger_tfs_dataframe
dframe.headers[
"TYPE"] = "TWISS" # MAD-X complains on TFS files with no "TYPE" header
write_location = tmp_path / "test.tfs"
write_tfs(write_location, dframe) # this will write an eol at eof
# The written TFS file should be accepted by MAD-X
with Madx() as madx:
madx.command.readtable(file=str(write_location),
table="test_table")
assert madx.table.test_table is not None # check table has loaded
# Check validity of the loaded table, here we use pandas.Series and assert_series_equal instead
# of numpy.array_equal to allow for (very) small relative numerical differences on loading
for column in dframe.columns:
assert column in madx.table.test_table
assert_series_equal(pandas.Series(
madx.table.test_table[column]),
dframe[column],
check_names=False)
def create_da_tfs(
jobname: str,
basedir: Path) -> Tuple[TfsDataFrame, TfsDataFrame, TfsDataFrame]:
""" Extracts data from db into dataframes, and writes and returns them."""
LOG.info("Gathering DA data into tfs-files.")
df_da = extract_da_data_from_db(jobname, basedir)
df_angle = _create_stats_df(df_da, ANGLE)
df_seed = _create_stats_df(df_da, SEED, global_index=0)
write_tfs(get_tfs_da_path(jobname, basedir), df_da)
write_tfs(get_tfs_da_angle_stats_path(jobname, basedir),
df_angle,
save_index=ANGLE)
write_tfs(get_tfs_da_seed_stats_path(jobname, basedir),
df_seed,
save_index=SEED)
return df_da, df_angle, df_seed
def test_fail_on_spaces_headers():
df = TfsDataFrame(headers={"allowed": 1, "not allowed": 2})
with pytest.raises(TfsFormatError):
write_tfs('', df)
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_messed_up_dataframe_fails_writes(
self, _messed_up_dataframe: TfsDataFrame):
messed_tfs = _messed_up_dataframe
with pytest.raises(ValueError):
write_tfs("", messed_tfs)
