diff --git a/pyerrors/input/pandas.py b/pyerrors/input/pandas.py index b3a2d3ef..67bd9bbd 100644 --- a/pyerrors/input/pandas.py +++ b/pyerrors/input/pandas.py @@ -2,6 +2,7 @@ import warnings import gzip import pandas as pd from ..obs import Obs +from ..correlators import Corr from .json import create_json_string, import_json_string @@ -23,7 +24,7 @@ def dump_df(df, fname, gz=True): out = df.copy() for column in out: - if isinstance(out[column][0], Obs): + if isinstance(out[column][0], (Obs, Corr)): out[column] = out[column].transform(lambda x: create_json_string(x, indent=0)) if not fname.endswith('.csv'): @@ -67,6 +68,6 @@ def load_df(fname, auto_gamma=False, gz=True): if re_import[column][0][:20] == '{"program":"pyerrors': re_import[column] = re_import[column].transform(lambda x: import_json_string(x, verbose=False)) if auto_gamma is True: - re_import[column].apply(Obs.gamma_method) + re_import[column].apply(lambda x: x.gamma_method()) return re_import diff --git a/tests/pandas_test.py b/tests/pandas_test.py index d4833656..658f4375 100644 --- a/tests/pandas_test.py +++ b/tests/pandas_test.py @@ -17,3 +17,14 @@ def test_df_export_import(tmp_path): pe.input.pandas.load_df((tmp_path / 'df_output.csv').as_posix(), gz=gz) +def test_df_Corr(tmp_path): + + my_corr = pe.Corr([pe.pseudo_Obs(-0.48, 0.04, "test"), pe.pseudo_Obs(-0.154, 0.03, "test")]) + + my_dict = {"int": 1, + "float": -0.01, + "Corr": my_corr} + my_df = pd.DataFrame([my_dict] * 5) + + pe.input.pandas.dump_df(my_df, (tmp_path / 'df_output').as_posix()) + reconstructed_df = pe.input.pandas.load_df((tmp_path / 'df_output').as_posix(), auto_gamma=True)