"""DataFrame class that is basically a wrapper around pandas.DataFrame which supports:
- Column aliases
- Units
- Case sensitivity
- search in the columns
- inputs and outputs into various additional formats: FITS, CSV/Ascii, ECSV
"""
import re
import sys
from collections.abc import Generator, Hashable, Iterable
from io import IOBase, TextIOWrapper
from os import PathLike
from typing import Any, Literal
import numpy.typing as npt
import pandas as pd
from pandas._typing import WriteBuffer
from . import convert
from .header import HeaderInfo
from .helpers import pretty_size_print
[docs]
class DataFrame(pd.DataFrame):
"""DataFrame class that is basically a wrapper around pandas.DataFrame which supports:
- Column aliases
- Units
- Case sensitivity
- search in the columns
- preserve attrs information at initialization from DataFrame
- inputs and outputs into various additional formats: FITS, CSV, ECSV
"""
[docs]
def __init__(
self,
data: npt.NDArray | npt.ArrayLike | Iterable | dict | pd.DataFrame,
info: HeaderInfo | None = None,
caseless: bool = False,
):
super().__init__(data)
attrs = getattr(data, "attrs", {})
self.attrs.update(attrs)
if info is not None:
self.attrs.update(info.header)
self.attrs["comments"] = info.comments
self.attrs["aliases"] = info.alias
self.attrs["units"] = info.units
self.attrs["caseless"] = caseless
self._clean_orphan_aliases()
@property
def header(self) -> HeaderInfo:
attrs = self.attrs.copy()
comments = attrs.pop("comments", {})
aliases = attrs.pop("aliases", {})
units = attrs.pop("units", {})
hdr = HeaderInfo(header=attrs, comments=comments, alias=aliases, units=units)
return hdr
@property
def nrows(self) -> int:
return len(self)
@property
def ncols(self) -> int:
return len(self.columns)
@property
def shape(self) -> tuple[int, int]:
return self.nrows, self.ncols
@property
def colnames(self) -> list[str]:
return list(self.columns)
@property
def nbytes(self):
"""number of bytes of the object"""
n = sum(
k.nbytes if hasattr(k, "nbytes") else sys.getsizeof(k)
for k in self.__dict__.values()
)
n += self.memory_usage(True, True).sum()
return n
[docs]
def resolve_alias(
self, colname: str | bytes | Hashable | Any
) -> str | bytes | Hashable | Any:
"""Resolve an alias to the actual column name
Parameters
----------
colname: str | bytes | Hashable | Any
Name or alias of the column
Returns
-------
str | bytes | Hashable | Any
Actual column name or initial value if not found
"""
if not self.attrs.get("caseless", False):
return self.attrs["aliases"].get(colname, colname)
else:
if isinstance(colname, (str, bytes)):
for k, v in self.attrs["aliases"].items():
if v.lower() == colname.lower():
return k
return colname
[docs]
def reverse_alias(
self, colname: str | bytes | Hashable | Any
) -> str | bytes | Hashable | Any:
"""Returns aliases of a given column
Parameters
----------
colname: str | bytes | Hashable | Any
Name or alias of the column
Returns
-------
str | bytes | Hashable | Any
Actual column name or initial value if not found
"""
_colname = self.resolve_alias(colname)
if _colname not in self.columns:
raise KeyError(f"Column '{_colname}' not found")
aliases = self.attrs.get("aliases", {})
return tuple([k for (k, v) in aliases.items() if (v == _colname)])
[docs]
def set_column_unit(self, colname: str, unit: str):
"""Set the unit for a column referenced by name or alias
Parameters
----------
colname: str
Name or alias of the column
unit: str
Unit to set for the column
"""
self.attrs["units"][self.resolve_alias(colname)] = str(unit)
def __getitem__(self, key: str | Any) -> Any:
"""Getitem accounting for aliases"""
return super().__getitem__(self.resolve_alias(key))
def __setitem__(self, key: str | Any, value: Any) -> None:
"""Setitem accounting for aliases"""
super().__setitem__(self.resolve_alias(key), value)
def _clean_orphan_aliases(self):
"""Make sure remaining aliases are correctly links to some data"""
aliases = self.attrs.get("aliases", {})
names = self.dtypes.index
self.attrs["aliases"] = {k: v for k, v in aliases.items() if v in names}
[docs]
def set_alias(self, alias: str, colname: str) -> None:
"""Set an alias for a column referenced by name or alias"""
if colname not in self.keys():
raise ValueError(f"Column '{colname}' does not exist")
self.attrs["aliases"][alias] = colname
[docs]
def keys( # pyright: ignore[reportIncompatibleMethodOverride]
self,
regexp: str | Iterable[str] | None = None,
full_match: bool = False,
skip_aliases: bool = False,
) -> list[str] | pd.Index:
"""
Return the data column names or a subset of it
Parameters
----------
regexp: str
pattern to filter the keys with
full_match: bool
if set, use :func:`re.fullmatch` instead of :func:`re.findall`
skip_aliases: bool
include aliases if unset
Try to apply the pattern at the start of the string, returning
a match object, or None if no match was found.
returns
-------
seq: sequence
sequence of keys
"""
if regexp in (None, "*"):
return super().keys()
if isinstance(regexp, (str, bytes)):
lbls = list(self.columns)
if not skip_aliases:
lbls.extend(self.attrs.get("aliases", {}).keys())
if regexp.count(",") > 0:
re_ = regexp.split(",")
elif regexp.strip().count(" ") > 0:
re_ = regexp.split(" ")
else:
re_ = [regexp]
fn = re.fullmatch if full_match else re.findall
matches = []
for subre in re_:
matches.extend([lbl for lbl in lbls if fn(subre, lbl)])
return matches
elif hasattr(regexp, "__iter__"):
matches = []
for rk in regexp:
matches.extend(
self.keys(rk, full_match=full_match, skip_aliases=skip_aliases)
)
return matches
else:
msg = f"Unexpected type {type(regexp)} for regexp"
raise ValueError(msg)
[docs]
def pandas_info(
self,
verbose: bool | None = None,
buf: WriteBuffer[str] | None = None,
max_cols: int | None = None,
memory_usage: bool | str | None = None,
show_counts: bool | None = None,
) -> None:
"""Print a concise summary of a DataFrame.
This method prints information about a DataFrame including
the index dtype and columns, non-null values and memory usage.
Parameters
----------
verbose : bool, optional
Whether to print the full summary. By default, the setting in
``pandas.options.display.max_info_columns`` is followed.
buf : writable buffer, defaults to sys.stdout
Where to send the output. By default, the output is printed to
sys.stdout. Pass a writable buffer if you need to further process
the output.
max_cols : int, optional
When to switch from the verbose to the truncated output. If the
DataFrame has more than `max_cols` columns, the truncated output
is used. By default, the setting in
``pandas.options.display.max_info_columns`` is used.
memory_usage : bool, str, optional
Specifies whether total memory usage of the DataFrame
elements (including the index) should be displayed. By default,
this follows the ``pandas.options.display.memory_usage`` setting.
True always show memory usage. False never shows memory usage.
A value of 'deep' is equivalent to "True with deep introspection".
Memory usage is shown in human-readable units (base-2
representation). Without deep introspection a memory estimation is
made based in column dtype and number of rows assuming values
consume the same memory amount for corresponding dtypes. With deep
memory introspection, a real memory usage calculation is performed
at the cost of computational resources. See the
:ref:`Frequently Asked Questions <df-memory-usage>` for more
details.
show_counts : bool, optional
Whether to show the non-null counts. By default, this is shown
only if the DataFrame is smaller than
``pandas.options.display.max_info_rows`` and
``pandas.options.display.max_info_columns``. A value of True always
shows the counts, and False never shows the counts.
Returns
-------
None
This method prints a summary of a DataFrame and returns None.
See Also
--------
DataFrame.describe: Generate descriptive statistics of DataFrame
columns.
DataFrame.memory_usage: Memory usage of DataFrame columns.
"""
return super().info(
verbose=verbose,
buf=buf,
max_cols=max_cols,
memory_usage=memory_usage,
show_counts=show_counts,
)
[docs]
def info(self, header: bool = False): # pyright: ignore
"""prints information on the table
Parameters
----------
header : bool, default False
If True, print the header information.
"""
txt = []
name = self.attrs.get("NAME", "Noname")
pprint_nbytes = pretty_size_print(self.nbytes)
txt.append(f"""{name:s}""")
txt.append(
f""" nrows={self.nrows:d}, ncols={self.ncols:d}, mem={pprint_nbytes:s}"""
)
# align keys properly
if header:
txt.append("Header:")
ignore = ("aliases", "units")
length = max([len(str(k)) for k in self.attrs.keys() if k not in ignore])
fmt = f"\t{{0:{length:d}s}} {{1}}"
for k, v in self.attrs.items():
if k not in ignore:
txt.append(fmt.format(k, str(v).strip()))
# display column descriptions properly aligned
units = self.attrs.get("units", {})
desc = self.attrs.get("comments", {})
values = [(k, units.get(k, ""), desc.get(k, "")) for k in self.colnames]
lengths = [[len(val) for val in colval] for colval in values]
lengths = list(map(max, (zip(*lengths))))
txt.append("Columns (name, units, description):")
fmt = "\t{{0:{0:d}s}} {{1:{1:d}s}} {{2:{2:d}s}}"
fmt = fmt.format(*(k + 1 for k in lengths))
for k, u, c in values:
txt.append(fmt.format(k, u, c))
aliases = self.attrs.get("aliases", {})
if aliases:
txt.append("Table contains alias(es):")
for k, v in aliases.items():
txt.append(f"\t{k:s} --> {v:s}")
print("\n".join(txt))
to_ascii = convert.ascii.to_ascii
to_csv = convert.ascii.to_csv # pyright: ignore
if convert.fits.from_fits is not None:
[docs]
@classmethod
def from_fits(cls, filename: str, extension_number: int = 1) -> "DataFrame":
"""Load a DataFrame from a FITS file.
Parameters
----------
filename : str
The path to the FITS file.
extension_number : int, optional
The extension number to load, by default 1.
Returns
-------
DataFrame
The loaded DataFrame.
"""
df, hdr = convert.fits.from_fits(filename, extension_number)
return cls(df, hdr)
if convert.fits.to_fits is not None:
[docs]
def to_fits(
self,
filename: str,
header_info: HeaderInfo | None = None,
output_verify: str = "exception",
checksum: bool = False,
index: bool = True,
overwrite: bool = False,
append: bool = False,
**kwargs,
) -> None:
"""Save a DataFrame to a FITS file.
Parameters
----------
filename : str
The path to the FITS file.
header_info : HeaderInfo | None, optional
Header information to save with the FITS file
by default None and taken from data.attrs
override data.attrs if provided
output_verify : str
Output verification option. Must be one of ``"fix"``, ``"silentfix"``,
``"ignore"``, ``"warn"``, or ``"exception"``. May also be any
combination of ``"fix"`` or ``"silentfix"`` with ``"+ignore"``,
``+warn``, or ``+exception" (e.g. ``"fix+warn"``). See :ref:`verify`
for more info.
checksum : bool, optional
If `True`, adds both ``DATASUM`` and ``CHECKSUM`` cards to the
headers of all HDU's written to the file
index : bool, optional
If `True`, includes the index in the FITS file. Default is `True`.
append : bool, optional
If `True`, appends the DataFrame to the FITS file. Default is `False`.
overwrite : bool, optional
If `True`, overwrites the DataFrame in the FITS file. Default is `False`.
**kwargs: dict
Additional keyword arguments to pass to the FITS writer.
"""
convert.fits.to_fits(
self,
filename,
header_info=self.header,
output_verify=output_verify,
checksum=checksum,
index=index,
overwrite=overwrite,
append=append,
**kwargs,
)
[docs]
@classmethod
def from_csv(
cls, filepath_or_buffer: str, *, commented_header: bool = False, **kwargs
) -> "DataFrame":
r"""
Read a CSV file into a DataFrame while preserving header information
Equivalent to `pd.read_csv` with preserved header information.
Also supports optionally iterating or breaking of the file
into chunks.
Additional help can be found in the online docs for
`IO Tools <https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html>`_.
Parameters
----------
filepath_or_buffer : str, path object or file-like object
Any valid string path is acceptable.
commented_header: bool, default False
Whether the column definition header line starts with a comment character.
commentchar: str, default '#'
Character to treat as a comment character.
sep : str, default ','
Character or regex pattern to treat as the delimiter. If ``sep=None``, the
C engine cannot automatically detect the separator, but the Python
parsing engine can, meaning the latter will be used and automatically
detect the separator from only the first valid row of the file by
Python's builtin sniffer tool, ``csv.Sniffer``.
In addition, separators longer than 1 character and different from
``'\s+'`` will be interpreted as regular expressions and will also force
the use of the Python parsing engine. Note that regex delimiters are prone
to ignoring quoted data. Regex example: ``'\r\t'``.
Returns
-------
DataFrame: pd.DataFrame
The parsed data as a pd.DataFrame.
header : HeaderInfo
The header information extracted from the file.
See also
--------
from_ascii: Read a ASCII file into a DataFrame
"""
return cls(
*convert.ascii.from_csv(
filepath_or_buffer, commented_header=commented_header, **kwargs
)
)
[docs]
@classmethod
def from_ascii(
cls, filepath_or_buffer: str, *, commented_header: bool = False, **kwargs
) -> "DataFrame":
"""Read an ASCII file into a DataFrame.
from_csv with delimiter set to " " by default
Parameters
----------
filepath_or_buffer : str, path object or file-like object
Any valid string path is acceptable. The string could be a URL. Valid
URL schemes include http, ftp, s3, and file. For file URLs, a host is
expected. A local file could be:
``file://localhost/path/to/table.csv``.
commented_header : bool, default False
Whether the header is commented or not.
**kwargs : dict
Additional keyword arguments passed to ``pd.read_csv``.
Returns
-------
DataFrame: pd.DataFrame
The parsed data as a pd.DataFrame.
header : HeaderInfo
The header information extracted from the file.
See also
--------
from_csv: Read a CSV file into a DataFrame
"""
return cls(
*convert.ascii.from_ascii(
filepath_or_buffer, commented_header=commented_header, **kwargs
)
)
[docs]
@classmethod
def from_ecsv(
cls, fname: str, **kwargs
) -> "DataFrame | Generator[DataFrame, Any, None]":
"""Read the content of an Enhanced Character Separated Values
Parameters
----------
fname: str
The name of the file to read.
Returns
-------
data: pd.DataFrame | Generator[pd.DataFrame, Any, None]
The data read from the file.
The chunked data read from the file.
"""
return cls(convert.ecsv.read(fname, **kwargs))
[docs]
def to_ecsv(
self,
df: pd.DataFrame,
fname: str | TextIOWrapper,
mode: str = "w",
**meta,
) -> None:
"""output data into ecsv file
Parameters
----------
df: pd.DataFrame
data to be written to the file.
fname: str
the name of the file to write.
mode: str
the mode to open the file.
meta: dict
meta data to be written to the header.
"""
convert.ecsv.write(df, fname, mode=mode, **meta)
if convert.hdf.from_hdf5 is not None:
@classmethod
def from_hdf5(
cls,
filename: str,
tablename: str | None = None,
*,
silent: bool = True,
**kwargs,
) -> "DataFrame":
"""Read a table from an HDF5 file.
Parameters
----------
filename: str
file to read from
tablename: str
node containing the table
silent: bool
skip verbose messages
Returns
-------
df: DataFrame
dataframe containing the table data
"""
data, hdr = convert.hdf.from_hdf5(
filename, tablename, silent=silent, **kwargs
)
return cls(data, hdr)
if convert.hdf.to_hdf5 is not None:
def to_hdf5(
self,
filename: str | PathLike,
*,
tablename: str | None = None,
mode: Literal["r", "w", "a", "r+"] = "w",
append: bool = False,
**kwargs,
) -> None:
"""
Write a pandas DataFrame to an HDF5 file.
Parameters
----------
filename : str or tables.File or PathLike
The filename or open HDF5 file to write to.
tablename : str, optional
The name of the table to write to.
mode : {'r', 'w', 'a', 'r+'}, default 'w'
The mode to open the file in.
append : bool, default False
Whether to append data to an existing file.
**kwargs
Additional keyword arguments to pass to tables.open_file.
Raises
------
Exception
If the HDF backend does not implement stream.
tables.FileModeError
If the file is already opened in a different mode.
ValueError
If something went wrong without much information from pytables.
"""
convert.hdf.to_hdf5(
self.data,
filename,
tablename=tablename,
header_info=self.header,
mode=mode,
append=append,
**kwargs,
)
if convert.votable.from_votable is not None:
[docs]
@classmethod
def from_votable(
cls,
fname: str | bytes | IOBase | PathLike,
*,
table_index: int = 0,
is_url: bool = False,
) -> "DataFrame":
"""Read a VOTable file and return a pandas DataFrame and header information.
Parameters
----------
fname : str, bytes, IOBase, PathLike
The filename or file-like object to read.
table_index : int, optional
The index of the table to read, by default 0.
is_url : bool, optional
Whether the file is a URL, by default False.
Returns
-------
df: DataFrame
Dataframe from VOTable file.
"""
df, hdr = convert.votable.from_votable(
fname, table_index=table_index, is_url=is_url
)
return cls(df, hdr)
[docs]
def testing():
filename = "/Users/fouesneau/Downloads/atlas9-munari.hires.grid.fits"
return DataFrame.from_fits(filename)