pyerrors.input.pandas

  1import gzip
  2import sqlite3
  3import warnings
  4from contextlib import closing
  5
  6import numpy as np
  7import pandas as pd
  8
  9from ..correlators import Corr
 10from ..obs import Obs
 11from .json import create_json_string, import_json_string
 12
 13
 14def to_sql(df, table_name, db, if_exists='fail', gz=True, **kwargs):
 15    """Write DataFrame including Obs or Corr valued columns to sqlite database.
 16
 17    Parameters
 18    ----------
 19    df : pandas.DataFrame
 20        Dataframe to be written to the database.
 21    table_name : str
 22        Name of the table in the database.
 23    db : str
 24        Path to the sqlite database.
 25    if exists : str
 26        How to behave if table already exists. Options 'fail', 'replace', 'append'.
 27    gz : bool
 28        If True the json strings are gzipped.
 29
 30    Returns
 31    -------
 32    None
 33    """
 34    se_df = _serialize_df(df, gz=gz)
 35    with closing(sqlite3.connect(db)) as con:
 36        se_df.to_sql(table_name, con=con, if_exists=if_exists, index=False, **kwargs)
 37
 38
 39def read_sql(sql, db, auto_gamma=False, **kwargs):
 40    """Execute SQL query on sqlite database and obtain DataFrame including Obs or Corr valued columns.
 41
 42    Parameters
 43    ----------
 44    sql : str
 45        SQL query to be executed.
 46    db : str
 47        Path to the sqlite database.
 48    auto_gamma : bool
 49        If True applies the gamma_method to all imported Obs objects with the default parameters for
 50        the error analysis. Default False.
 51
 52    Returns
 53    -------
 54    data : pandas.DataFrame
 55        Dataframe with the content of the sqlite database.
 56    """
 57    with closing(sqlite3.connect(db)) as con:
 58        extract_df = pd.read_sql(sql, con=con, **kwargs)
 59    return _deserialize_df(extract_df, auto_gamma=auto_gamma)
 60
 61
 62def dump_df(df, fname, gz=True):
 63    """Exports a pandas DataFrame containing Obs valued columns to a (gzipped) csv file.
 64
 65    Before making use of pandas to_csv functionality Obs objects are serialized via the standardized
 66    json format of pyerrors.
 67
 68    Parameters
 69    ----------
 70    df : pandas.DataFrame
 71        Dataframe to be dumped to a file.
 72    fname : str
 73        Filename of the output file.
 74    gz : bool
 75        If True, the output is a gzipped csv file. If False, the output is a csv file.
 76
 77    Returns
 78    -------
 79    None
 80    """
 81    for column in df:
 82        serialize = _need_to_serialize(df[column])
 83        if not serialize:
 84            if all(isinstance(entry, (int, np.integer, float, np.floating)) for entry in df[column]):
 85                if any([np.isnan(entry) for entry in df[column]]):
 86                    warnings.warn("nan value in column " + column + " will be replaced by None", UserWarning, stacklevel=2)
 87
 88    out = _serialize_df(df, gz=False)
 89
 90    if not fname.endswith('.csv'):
 91        fname += '.csv'
 92
 93    if gz is True:
 94        if not fname.endswith('.gz'):
 95            fname += '.gz'
 96        out.to_csv(fname, index=False, compression='gzip')
 97    else:
 98        out.to_csv(fname, index=False)
 99
100
101def load_df(fname, auto_gamma=False, gz=True):
102    """Imports a pandas DataFrame from a csv.(gz) file in which Obs objects are serialized as json strings.
103
104    Parameters
105    ----------
106    fname : str
107        Filename of the input file.
108    auto_gamma : bool
109        If True applies the gamma_method to all imported Obs objects with the default parameters for
110        the error analysis. Default False.
111    gz : bool
112        If True, assumes that data is gzipped. If False, assumes JSON file.
113
114    Returns
115    -------
116    data : pandas.DataFrame
117        Dataframe with the content of the sqlite database.
118    """
119    if not fname.endswith('.csv') and not fname.endswith('.gz'):
120        fname += '.csv'
121
122    if gz is True:
123        if not fname.endswith('.gz'):
124            fname += '.gz'
125        with gzip.open(fname) as f:
126            re_import = pd.read_csv(f, keep_default_na=False)
127    else:
128        if fname.endswith('.gz'):
129            warnings.warn(f"Trying to read from {fname} without unzipping!", UserWarning, stacklevel=2)
130        re_import = pd.read_csv(fname, keep_default_na=False)
131
132    return _deserialize_df(re_import, auto_gamma=auto_gamma)
133
134
135def _serialize_df(df, gz=False):
136    """Serializes all Obs or Corr valued columns into json strings according to the pyerrors json specification.
137
138    Parameters
139    ----------
140    df : pandas.DataFrame
141        DataFrame to be serilized.
142    gz: bool
143        gzip the json string representation. Default False.
144    """
145    out = df.copy()
146    for column in out:
147        serialize = _need_to_serialize(out[column])
148
149        if serialize is True:
150            out[column] = out[column].transform(lambda x: create_json_string(x, indent=0) if not _is_null(x) else None)
151            if gz is True:
152                out[column] = out[column].transform(lambda x: gzip.compress(x.encode('utf-8')) if not _is_null(x) else gzip.compress(b''))
153    return out
154
155
156def _deserialize_df(df, auto_gamma=False):
157    """Deserializes all pyerrors json strings into Obs or Corr objects according to the pyerrors json specification.
158
159    Parameters
160    ----------
161    df : pandas.DataFrame
162        DataFrame to be deserilized.
163    auto_gamma : bool
164        If True applies the gamma_method to all imported Obs objects with the default parameters for
165        the error analysis. Default False.
166
167    Notes:
168    ------
169    In case any column of the DataFrame is gzipped it is gunzipped in the process.
170    """
171    # In pandas 3+, string columns use 'str' dtype instead of 'object'
172    string_like_dtypes = ["object", "str"] if int(pd.__version__.split(".")[0]) >= 3 else ["object"]
173    for column in df.select_dtypes(include=string_like_dtypes):
174        if len(df[column]) == 0:
175            continue
176        if isinstance(df[column].iloc[0], bytes):
177            if df[column].iloc[0].startswith(b"\x1f\x8b\x08\x00"):
178                df[column] = df[column].transform(lambda x: gzip.decompress(x).decode('utf-8') if not pd.isna(x) else '')
179
180        if df[column].notna().any():
181            df[column] = df[column].replace({r'^$': None}, regex=True)
182            i = 0
183            while i < len(df[column]) and pd.isna(df[column].iloc[i]):
184                i += 1
185            if i < len(df[column]) and isinstance(df[column].iloc[i], str):
186                if '"program":' in df[column].iloc[i][:20]:
187                    df[column] = df[column].transform(lambda x: import_json_string(x, verbose=False) if not pd.isna(x) else None)
188                    if auto_gamma is True:
189                        if isinstance(df[column].iloc[i], list):
190                            df[column].apply(lambda x: [o.gm() if o is not None else x for o in x] if x is not None else x)
191                        else:
192                            df[column].apply(lambda x: x.gm() if x is not None else x)
193        # Convert NA values back to Python None for compatibility with `x is None` checks
194        if df[column].isna().any():
195            df[column] = df[column].astype(object).where(df[column].notna(), None)
196    return df
197
198
199def _need_to_serialize(col):
200    serialize = False
201    i = 0
202    while i < len(col) and _is_null(col.iloc[i]):
203        i += 1
204    if i == len(col):
205        return serialize
206    if isinstance(col.iloc[i], (Obs, Corr)):
207        serialize = True
208    elif isinstance(col.iloc[i], list):
209        if all(isinstance(o, Obs) for o in col.iloc[i]):
210            serialize = True
211    return serialize
212
213
214def _is_null(val):
215    """Check if a value is null (None or NA), handling list/array values."""
216    return False if isinstance(val, (list, np.ndarray)) else pd.isna(val)