"""Typed YAML I/O for retrieval folders.
A retrieval folder lives at ``Files/Targets/<target>/Retrievals/<run>/`` and
historically held two CSV files:
- ``df_general_info.csv`` — 2-column ``Variable``/``Value`` key/value table
- ``df_parameters.csv`` — 15-column DataFrame, one row per fitted parameter
Both files round-tripped via :func:`pandas.DataFrame.to_csv`/``read_csv`` lose
all dtype information: every value comes back as ``object``. This module
replaces that with typed YAML persistence while keeping legacy CSVs readable
through a lazy conversion path. Writes go through atomic temp + ``os.replace``
so a mid-write crash never leaves a partial file.
Public API
----------
- :func:`read_general_info(folder)` -> typed ``dict``
- :func:`read_general_info_df(folder)` -> 2-col legacy-shaped DataFrame
- :func:`write_general_info(folder, data)` -> ``Path`` (accepts dict or DF)
- :func:`read_parameters(folder)` -> 15-col DataFrame, NaN preserved
- :func:`write_parameters(folder, df)` -> ``Path``
- :func:`ensure_yaml(folder)` -> ``{"general_info": Path, "parameters": Path}``
- :class:`RetrievalIOError`
The CSV is **never** deleted or renamed by this module.
"""
from __future__ import annotations
import datetime as _dt
import os
from pathlib import Path
from typing import Any, Dict, Mapping, Sequence, Union
import numpy as np
import pandas as pd
import yaml
from GUIBRUSHR.core.io._coerce import (
coerce_bool as _coerce_bool,
coerce_list as _coerce_list,
coerce_scalar as _coerce_scalar,
is_nullish as _is_nullish,
)
# Canonical 15-column schema for df_parameters. Mirrors
# Constant_Variables.COLUMN_DF_PARAMETERS but is duplicated here to avoid
# a runtime import of GUI-side constants from this low-level module.
PARAMETER_COLUMNS: list[str] = [
"name",
"is_present",
"molec",
"value",
"scale",
"range_min",
"range_max",
"rayleigh_species",
"in_bestpars",
"mass",
"sigma_prior",
"molec_formula",
"constant_vmr",
"isotope",
"opacity_name_lr",
]
# Columns whose legacy CSV representation is "1"/"0" or "True"/"False"
# but the canonical typed value is a real bool.
_BOOL_PARAMETER_COLUMNS = {
"is_present", "molec", "rayleigh_species", "in_bestpars", "constant_vmr",
} # match Constant_Variables.TYPES_DF_PARAMETERS
GENERAL_INFO_FILENAME_CSV = "df_general_info.csv"
GENERAL_INFO_FILENAME_YAML = "df_general_info.yaml"
PARAMETERS_FILENAME_CSV = "df_parameters.csv"
PARAMETERS_FILENAME_YAML = "df_parameters.yaml"
PathLike = Union[str, os.PathLike]
[docs]
class RetrievalIOError(Exception):
"""Raised on retrieval-folder I/O errors (empty YAML, missing file, schema mismatch)."""
# ---------------------------------------------------------------------------
# YAML representers — emit native scalars from numpy and force | block style
# for any string containing a newline. Registered once at import time.
# ---------------------------------------------------------------------------
def _represent_numpy_int(dumper: yaml.Dumper, data: np.integer) -> yaml.ScalarNode:
return dumper.represent_int(int(data.item()))
def _represent_numpy_float(dumper: yaml.Dumper, data: np.floating) -> yaml.ScalarNode:
return dumper.represent_float(float(data.item()))
def _represent_numpy_bool(dumper: yaml.Dumper, data: np.bool_) -> yaml.ScalarNode:
return dumper.represent_bool(bool(data.item()))
def _represent_numpy_array(dumper: yaml.Dumper, data: np.ndarray):
return dumper.represent_list(data.tolist())
def _represent_str(dumper: yaml.Dumper, data: str):
if "\n" in data:
return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
return dumper.represent_scalar("tag:yaml.org,2002:str", data)
for _np_int_t in (np.int8, np.int16, np.int32, np.int64,
np.uint8, np.uint16, np.uint32, np.uint64):
yaml.SafeDumper.add_representer(_np_int_t, _represent_numpy_int)
for _np_float_t in (np.float16, np.float32, np.float64):
yaml.SafeDumper.add_representer(_np_float_t, _represent_numpy_float)
yaml.SafeDumper.add_representer(np.bool_, _represent_numpy_bool)
yaml.SafeDumper.add_representer(np.ndarray, _represent_numpy_array)
yaml.SafeDumper.add_representer(str, _represent_str)
# ---------------------------------------------------------------------------
# Atomic write
# ---------------------------------------------------------------------------
def _atomic_write(path: Path, content: str) -> None:
"""Write ``content`` to ``path`` atomically (tmp file + fsync + ``os.replace``)."""
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
tmp = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")
try:
with open(tmp, "w", encoding="utf-8") as fp:
fp.write(content)
fp.flush()
os.fsync(fp.fileno())
os.replace(tmp, path)
except BaseException:
# On any failure (including KeyboardInterrupt) try to remove the tmp
# so we don't leak orphaned ``*.tmp.<pid>`` files in retrieval folders.
try:
tmp.unlink()
except FileNotFoundError:
pass
raise
def _today_iso() -> str:
return _dt.date.today().isoformat()
def _provenance_header(source: str | None) -> str:
if source is None:
return f"# generated {_today_iso()}\n"
return f"# generated from {source} on {_today_iso()}\n"
def _dump_yaml(data: Any) -> str:
return yaml.safe_dump(data, sort_keys=False, allow_unicode=True, default_flow_style=False)
def _safe_load_yaml(path: Path) -> Any:
with open(path, "r", encoding="utf-8") as fp:
return yaml.safe_load(fp)
# ---------------------------------------------------------------------------
# General info
# ---------------------------------------------------------------------------
# Heuristic keys whose value is a list (when present). All current legacy
# stringified-list keys we have observed; the readers also fall back to
# parenthesis/bracket detection in _coerce_scalar's caller to handle anything
# that *looks* like a list at runtime.
_KNOWN_LIST_KEYS = {"instruments_hr_list", "instruments_lr_list"}
def _typed_value_for_general_info(key: str, raw: Any) -> Any:
"""Coerce a single legacy general-info value to a typed Python value.
Bracket-detection is restricted to ``_KNOWN_LIST_KEYS`` so user-typed
values like ``comments = "[see appendix]"`` are not silently coerced
into single-element lists.
"""
if _is_nullish(raw):
return None
if isinstance(raw, (list, tuple, np.ndarray)):
return _coerce_list(raw)
if key in _KNOWN_LIST_KEYS:
return _coerce_list(raw)
return _coerce_scalar(raw)
def _general_info_dict_from_csv_df(df: pd.DataFrame) -> Dict[str, Any]:
"""Convert the 2-col ``Variable``/``Value`` legacy DF into a typed dict."""
if "Variable" not in df.columns or "Value" not in df.columns:
raise RetrievalIOError(
"df_general_info CSV missing 'Variable'/'Value' columns; "
f"got {list(df.columns)!r}"
)
out: Dict[str, Any] = {}
for _, row in df.iterrows():
key = row["Variable"]
if _is_nullish(key):
continue
out[str(key)] = _typed_value_for_general_info(str(key), row["Value"])
return out
def _read_general_info_csv(csv_path: Path) -> Dict[str, Any]:
df = pd.read_csv(csv_path)
# to_csv was called without index=False, so the first col may be an unnamed index.
drop_cols = [c for c in df.columns if str(c).startswith("Unnamed:")]
if drop_cols:
df = df.drop(columns=drop_cols)
return _general_info_dict_from_csv_df(df)
def _normalize_general_info_input(data: Mapping[str, Any] | pd.DataFrame) -> Dict[str, Any]:
"""Accept either a typed dict or a 2-col legacy DF and return a typed dict."""
if isinstance(data, pd.DataFrame):
return _general_info_dict_from_csv_df(data)
if isinstance(data, Mapping):
return {str(k): v for k, v in data.items()}
raise RetrievalIOError(
f"write_general_info: expected dict or 2-col DataFrame, got {type(data).__name__}"
)
[docs]
def read_general_info(folder: PathLike) -> Dict[str, Any]:
"""Return the general-info table as a typed ``dict``.
Prefers ``df_general_info.yaml``; falls back to legacy CSV and lazily
writes a YAML sibling so subsequent reads are fast. The CSV is never
deleted or modified.
Raises :class:`RetrievalIOError` if neither file exists or the YAML is
empty.
"""
folder = Path(folder)
yaml_path = folder / GENERAL_INFO_FILENAME_YAML
csv_path = folder / GENERAL_INFO_FILENAME_CSV
if yaml_path.exists():
loaded = _safe_load_yaml(yaml_path)
if loaded is None:
raise RetrievalIOError(f"{yaml_path} is empty")
if not isinstance(loaded, Mapping):
raise RetrievalIOError(
f"{yaml_path} is not a mapping (got {type(loaded).__name__})"
)
return dict(loaded)
if csv_path.exists():
data = _read_general_info_csv(csv_path)
# lazy-convert
try:
_write_general_info_yaml(yaml_path, data, source=GENERAL_INFO_FILENAME_CSV)
except OSError:
# read-only filesystem etc — still return the data
pass
return data
raise RetrievalIOError(
f"No df_general_info.yaml or df_general_info.csv in folder {folder}"
)
[docs]
def read_general_info_df(folder: PathLike) -> pd.DataFrame:
"""Return the general-info table as a 2-col legacy-shaped DataFrame.
Used by call sites that pass the result to :func:`get_csv_value`. List
values are kept as native Python lists in the ``Value`` column (object dtype).
"""
data = read_general_info(folder)
rows = [(k, v) for k, v in data.items()]
return pd.DataFrame(rows, columns=["Variable", "Value"])
def _write_general_info_yaml(path: Path, data: Mapping[str, Any], *, source: str | None) -> Path:
header = _provenance_header(source)
body = _dump_yaml(dict(data))
_atomic_write(path, header + body)
return path
[docs]
def write_general_info(folder: PathLike, data: Mapping[str, Any] | pd.DataFrame) -> Path:
"""Write ``df_general_info.yaml`` with typed values. Accepts dict or 2-col DF."""
folder = Path(folder)
typed = _normalize_general_info_input(data)
return _write_general_info_yaml(
folder / GENERAL_INFO_FILENAME_YAML, typed, source=None
)
# ---------------------------------------------------------------------------
# Parameters
# ---------------------------------------------------------------------------
def _typed_value_for_param(col: str, raw: Any) -> Any:
"""Coerce a single legacy parameters cell to its typed Python value."""
if col in _BOOL_PARAMETER_COLUMNS:
if _is_nullish(raw):
return False
try:
return _coerce_bool(raw)
except ValueError:
# fall through to scalar coercion if oddly shaped
return _coerce_scalar(raw)
return _coerce_scalar(raw)
def _params_records_from_csv_df(df: pd.DataFrame) -> list[dict]:
"""Convert a parameters DataFrame loaded from CSV into typed records.
Permissive on missing columns: schema fields added in later versions
(e.g. ``opacity_name_lr``) are absent from older retrievals — fill with
NaN rather than failing. The required-name check stays strict.
"""
drop_cols = [c for c in df.columns if str(c).startswith("Unnamed:")]
if drop_cols:
df = df.drop(columns=drop_cols)
if "name" not in df.columns:
raise RetrievalIOError(
f"df_parameters CSV missing required 'name' column; got {list(df.columns)!r}"
)
for col in PARAMETER_COLUMNS:
if col not in df.columns:
df[col] = np.nan
df = df[PARAMETER_COLUMNS]
records: list[dict] = []
for _, row in df.iterrows():
rec: dict = {}
for col in PARAMETER_COLUMNS:
typed = _typed_value_for_param(col, row[col])
if isinstance(typed, float) and np.isnan(typed):
continue # NaN cells omitted from YAML
if typed is None and col not in _BOOL_PARAMETER_COLUMNS:
continue
rec[col] = typed
records.append(rec)
return records
def _df_from_records(records: Sequence[Mapping[str, Any]]) -> pd.DataFrame:
"""Rebuild the 15-col DataFrame, restoring NaN for omitted keys."""
full: list[dict] = []
for rec in records:
row = {col: rec.get(col, np.nan) for col in PARAMETER_COLUMNS}
for bool_col in _BOOL_PARAMETER_COLUMNS:
if bool_col in rec:
row[bool_col] = bool(rec[bool_col])
full.append(row)
return pd.DataFrame(full, columns=PARAMETER_COLUMNS)
def _records_from_df(df: pd.DataFrame) -> list[dict]:
"""Convert the in-memory parameters DataFrame to YAML-ready records."""
missing = [c for c in PARAMETER_COLUMNS if c not in df.columns]
if missing:
raise RetrievalIOError(
f"write_parameters: DataFrame missing columns {missing}"
)
df = df[PARAMETER_COLUMNS]
records: list[dict] = []
for _, row in df.iterrows():
rec: dict = {}
for col in PARAMETER_COLUMNS:
value = row[col]
if col in _BOOL_PARAMETER_COLUMNS:
if _is_nullish(value):
rec[col] = False
else:
try:
rec[col] = _coerce_bool(value)
except ValueError:
rec[col] = False
continue
if _is_nullish(value):
continue
if isinstance(value, np.generic):
value = value.item()
rec[col] = value
records.append(rec)
return records
[docs]
def read_parameters(folder: PathLike) -> pd.DataFrame:
"""Return the 15-col parameters DataFrame, NaN preserved for missing cells."""
folder = Path(folder)
yaml_path = folder / PARAMETERS_FILENAME_YAML
csv_path = folder / PARAMETERS_FILENAME_CSV
if yaml_path.exists():
loaded = _safe_load_yaml(yaml_path)
if loaded is None:
raise RetrievalIOError(f"{yaml_path} is empty")
if not isinstance(loaded, Mapping) or "parameters" not in loaded:
raise RetrievalIOError(
f"{yaml_path} missing top-level 'parameters' key"
)
records = loaded["parameters"] or []
return _df_from_records(records)
if csv_path.exists():
df_csv = pd.read_csv(csv_path)
records = _params_records_from_csv_df(df_csv)
try:
_write_parameters_yaml(yaml_path, records, source=PARAMETERS_FILENAME_CSV)
except OSError:
pass
return _df_from_records(records)
raise RetrievalIOError(
f"No df_parameters.yaml or df_parameters.csv in folder {folder}"
)
def _write_parameters_yaml(path: Path, records: Sequence[Mapping[str, Any]], *, source: str | None) -> Path:
header = _provenance_header(source)
body = _dump_yaml({"parameters": [dict(r) for r in records]})
_atomic_write(path, header + body)
return path
[docs]
def write_parameters(folder: PathLike, df: pd.DataFrame) -> Path:
"""Write ``df_parameters.yaml``. NaN cells are omitted from the YAML."""
folder = Path(folder)
if not isinstance(df, pd.DataFrame):
raise RetrievalIOError(
f"write_parameters: expected DataFrame, got {type(df).__name__}"
)
records = _records_from_df(df)
return _write_parameters_yaml(
folder / PARAMETERS_FILENAME_YAML, records, source=None
)
# ---------------------------------------------------------------------------
# ensure_yaml — lazy CSV -> YAML migration. Idempotent. Never deletes CSV.
# ---------------------------------------------------------------------------
[docs]
def ensure_yaml(folder: PathLike) -> Dict[str, Path]:
"""Ensure YAML siblings exist for any legacy CSV in ``folder``.
Returns a mapping ``{"general_info": <path>, "parameters": <path>}`` for
each kind that exists (in either format). The CSV is left untouched.
Calling twice is a no-op (idempotent).
"""
folder = Path(folder)
out: Dict[str, Path] = {}
yaml_g = folder / GENERAL_INFO_FILENAME_YAML
csv_g = folder / GENERAL_INFO_FILENAME_CSV
if yaml_g.exists():
out["general_info"] = yaml_g
elif csv_g.exists():
data = _read_general_info_csv(csv_g)
_write_general_info_yaml(yaml_g, data, source=GENERAL_INFO_FILENAME_CSV)
out["general_info"] = yaml_g
yaml_p = folder / PARAMETERS_FILENAME_YAML
csv_p = folder / PARAMETERS_FILENAME_CSV
if yaml_p.exists():
out["parameters"] = yaml_p
elif csv_p.exists():
df_csv = pd.read_csv(csv_p)
records = _params_records_from_csv_df(df_csv)
_write_parameters_yaml(yaml_p, records, source=PARAMETERS_FILENAME_CSV)
out["parameters"] = yaml_p
return out