eng-lab/hwe/testdata.py

"""Read measured data (and datasheet tables) with their units attached.

Every CSV this engine reads names each column's unit in its header, the way
phyphox exports do:  Time (s), Acceleration z (m/s^2)
Square brackets work too:  thrust [gf]
A column without a unit is refused, because a measurement without a unit is
not a measurement (EL01, section 2.2).

read_csv returns a Table: an ordered dict from column name to a Quantity that
wraps a numpy array, plus the file it came from.
"""
from __future__ import annotations

import csv
import re
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

from hwe.units import Q_, ureg

_HEADER = re.compile(r"^\s*(?P<name>.*?)\s*[\(\[](?P<unit>[^\)\]]*)[\)\]]\s*$")

# Header spellings that pint does not read directly.
_UNIT_ALIASES = {
    "%": "percent", "°C": "degC", "deg": "degree", "g/W": "gf/W",
    "": "dimensionless", "-": "dimensionless",
}


@dataclass
class Table:
    path: Path
    columns: dict = field(default_factory=dict)      # name -> Quantity array

    def __getitem__(self, name):
        return self.columns[name]

    def names(self):
        return list(self.columns)

    def find(self, *words):
        """The first column whose name contains every word (case-insensitive)."""
        for name, values in self.columns.items():
            low = name.lower()
            if all(w.lower() in low for w in words):
                return values
        raise KeyError(f"no column containing {words} in {self.path.name}; have {self.names()}")


def parse_header(header: str) -> tuple[str, str]:
    """'Acceleration z (m/s^2)' -> ('Acceleration z', 'm/s^2')"""
    match = _HEADER.match(header)
    if not match:
        raise ValueError(f"column {header!r} has no unit; write it as 'name (unit)'")
    name, unit = match.group("name"), match.group("unit").strip()
    unit = _UNIT_ALIASES.get(unit, unit)
    ureg.parse_units(unit)            # fail now, with the column named, if pint can't read it
    return name, unit


def read_csv(path: Path | str) -> Table:
    """Read a CSV whose headers carry units. Decimal points, comma separators."""
    path = Path(path)
    with open(path, newline="", encoding="utf-8-sig") as f:
        sample = f.read(4096)
        f.seek(0)
        dialect = csv.Sniffer().sniff(sample, delimiters=",;\t")
        rows = list(csv.reader(f, dialect))
    header, body = rows[0], [r for r in rows[1:] if any(cell.strip() for cell in r)]
    table = Table(path=path)
    for j, head in enumerate(header):
        name, unit = parse_header(head)
        numbers = np.array([float(r[j]) for r in body])
        table.columns[name] = Q_(numbers, unit)
    return table