From 3ed8d5efb0a050f0905165556f19a585675ee708 Mon Sep 17 00:00:00 2001 From: Jacob Wilkins <46597752+oerc0122@users.noreply.github.com> Date: Thu, 1 Aug 2024 15:07:51 +0100 Subject: [PATCH] Add table writer to utils for standard format tables (#229) * Add table writer to utils for standard format tables * Update conf.py * Respond to review --------- Co-authored-by: Alin Marin Elena --- docs/source/conf.py | 1 + janus_core/helpers/utils.py | 186 +++++++++++++++++++++++++++++++++++- 2 files changed, 186 insertions(+), 1 deletion(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 8359d3d3..2b992253 100755 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -204,4 +204,5 @@ ("py:class", "phonopy.structure.atoms.PhonopyAtoms"), ("py:class", "ase.optimize.optimize.Optimizer"), ("py:class", "codecarbon.emissions_tracker.OfflineEmissionsTracker"), + ("py:class", "_io.StringIO"), ] diff --git a/janus_core/helpers/utils.py b/janus_core/helpers/utils.py index 8a7fb0f2..d54b93c8 100644 --- a/janus_core/helpers/utils.py +++ b/janus_core/helpers/utils.py @@ -2,8 +2,9 @@ from abc import ABC from collections.abc import Collection, Sequence +from io import StringIO from pathlib import Path -from typing import Optional, get_args +from typing import Any, Literal, Optional, TextIO, get_args from ase import Atoms from ase.io import write @@ -337,3 +338,186 @@ def output_structs( if write_results: write_kwargs.setdefault("write_results", not invalidate_calc) write(images=images, **write_kwargs) + + +def write_table( + fmt: Literal["ascii", "csv"], + file: Optional[TextIO] = None, + units: Optional[dict[str, str]] = None, + formats: Optional[dict[str, str]] = None, + **columns, +) -> Optional[StringIO]: + """ + Dump a table in a standard format. + + Table columns are passed as kwargs, with the column header being + the keyword name and data the value. + + Each header also supports an optional "
_units" or + "
_format" kwarg to define units and format for the column. + These can also be passed explicitly through the respective + dictionaries where the key is the "header". + + Parameters + ---------- + fmt : {'ascii', 'csv'} + Format to write table in. + file : TextIO + File to dump to. If unspecified function returns + io.StringIO object simulating file. + units : dict[str, str] + Units as ``{key: unit}``: + + key + To align with those in `columns`. + unit + String defining unit. + + Units which do not match any columns in `columns` are + ignored. + formats : dict[str, str] + Output formats as ``{key: format}``: + + key + To align with those in `columns`. + format + Python magic format string to use. + **columns : dict[str, Any] + Dictionary of columns names to data with optional + "
_units"/"
_format" to define units/format. + + See: Examples. + + Returns + ------- + Optional[StringIO] + If no file given write columns to StringIO. + + Notes + ----- + Passing "kwarg_units" or "kwarg_format" takes priority over + those passed in the `units`/`formats` dicts. + + Examples + -------- + >>> data = write_table(fmt="ascii", single=(1, 2, 3), + ... double=(2,4,6), double_units="THz") + >>> print(*data, sep="", end="") + # single | double [THz] + 1 2 + 2 4 + 3 6 + >>> data = write_table(fmt="csv", a=(3., 5.), b=(11., 12.), + ... units={'a': 'eV', 'b': 'Ha'}) + >>> print(*data, sep="", end="") + a [eV],b [Ha] + 3.0,11.0 + 5.0,12.0 + >>> data = write_table(fmt="csv", a=(3., 5.), b=(11., 12.), + ... formats={"a": ".3f"}) + >>> print(*data, sep="", end="") + a,b + 3.000,11.0 + 5.000,12.0 + """ + units = units if units else {} + units.update( + { + key.removesuffix("_units"): val + for key, val in columns.items() + if key.endswith("_units") + } + ) + + formats = formats if formats else {} + formats.update( + { + key.removesuffix("_format"): val + for key, val in columns.items() + if key.endswith("_format") + } + ) + + columns = {key: val for key, val in columns.items() if not key.endswith("_units")} + + header = [ + f"{datum}" + (f" [{units.get(datum)}]" if datum in units else "") + for datum in columns + ] + + dump_loc = file if file is not None else StringIO() + + write_fmt = [formats.get(key, "") for key in columns] + + if fmt == "ascii": + _dump_ascii(dump_loc, header, columns, write_fmt) + elif fmt == "csv": + _dump_csv(dump_loc, header, columns, write_fmt) + + if file is None: + dump_loc.seek(0) + return dump_loc + return None + + +def _dump_ascii( + file: TextIO, + header: list[str], + columns: dict[str, Sequence[Any]], + formats: Sequence[str], +): + """ + Dump data as an ASCII style table. + + Parameters + ---------- + file : TextIO + File to dump to. + header : list[str] + Column name information. + columns : dict[str, Sequence[Any]] + Column data by key (ordered with header info). + formats : Sequence[str] + Python magic string formats to apply + (must align with header info). + + See Also + -------- + write_table : Main entry function. + """ + + print(f"# {' | '.join(header)}", file=file) + + for cols in zip(*columns.values()): + print(*map(format, cols, formats), file=file) + + +def _dump_csv( + file: TextIO, + header: list[str], + columns: dict[str, Sequence[Any]], + formats: Sequence[str], +): + """ + Dump data as a csv style table. + + Parameters + ---------- + file : TextIO + File to dump to. + header : list[str] + Column name information. + columns : dict[str, Sequence[Any]] + Column data by key (ordered with header info). + formats : Sequence[str] + Python magic string formats to apply + (must align with header info). + + See Also + -------- + write_table : Main entry function. + """ + print(",".join(header), file=file) + + for cols in zip(*columns.values()): + print(",".join(map(format, cols, formats)), file=file)