eemkit 0.1.0__py3-none-any.whl

eemkit/__init__.py

@@ -0,0 +1,51 @@
+"""eemkit: EEM processing and PARAFAC modelling of dissolved organic matter fluorescence."""
+
+__version__ = "0.1.0"
+
+from eemkit import corrections, indices, parafac, pipeline
+from eemkit.core.dataset import stack
+from eemkit.core.grid import infer_common_grid, regular_grid, resample
+from eemkit.core.history import get_history
+from eemkit.errors import (
+    Cancelled,
+    CorrectionError,
+    EEMKitError,
+    EEMKitWarning,
+    FileFormatError,
+    GridError,
+    PairingError,
+    SettingsError,
+)
+from eemkit.io.cary60 import mean_cycles, read_cary60
+from eemkit.io.cary_eclipse import read_cary_eclipse
+from eemkit.io.discover import scan_folder, suggest_pairing
+from eemkit.io.loader import load_dataset
+from eemkit.io.pairing import read_pairing
+
+__all__ = [
+    "Cancelled",
+    "CorrectionError",
+    "EEMKitError",
+    "EEMKitWarning",
+    "FileFormatError",
+    "GridError",
+    "PairingError",
+    "SettingsError",
+    "__version__",
+    "corrections",
+    "get_history",
+    "indices",
+    "infer_common_grid",
+    "load_dataset",
+    "mean_cycles",
+    "parafac",
+    "pipeline",
+    "read_cary60",
+    "read_cary_eclipse",
+    "read_pairing",
+    "regular_grid",
+    "resample",
+    "scan_folder",
+    "stack",
+    "suggest_pairing",
+]

eemkit/absorbance.py

@@ -0,0 +1,216 @@
+"""Chromophoric DOM (CDOM) metrics from absorbance spectra.
+
+=================  ===================================================  ==========================
+metric             definition (defaults)                                reference
+=================  ===================================================  ==========================
+``a254``, ``a350``  Napierian absorption coefficient a = ln(10) A / l    Green & Blough 1994
+                   (m^-1, path length l in m)
+``S275_295``       slope of ln a(lambda) against lambda, 275-295 nm      Helms et al. 2008
+``S350_400``       the same over 350-400 nm (nm^-1, reported positive)  Helms et al. 2008
+``SR``             S275_295 / S350_400                                  Helms et al. 2008
+``E2_E3``          A(250) / A(365)                                      De Haan & De Boer 1987
+``SUVA254``        A(254) / l (m^-1, decadic) / DOC (mg L^-1)           Weishaar et al. 2003
+=================  ===================================================  ==========================
+
+Before any metric, each spectrum is corrected for its baseline offset by
+subtracting its mean absorbance over ``baseline`` (default 680-700 nm, where
+CDOM does not absorb; Green & Blough 1994, Helms et al. 2008). The absorbance
+must already be blank-corrected
+(:func:`~eemkit.corrections.subtract_absorbance_blank`).
+
+Low signal is never turned into a number: a metric that would need a
+non-positive absorbance, or a slope whose log-linear fit has R^2 below
+``min_r2``, is NaN, and the reason is written in the ``cdom_notes`` column.
+"""
+
+from __future__ import annotations
+
+import warnings
+from collections.abc import Sequence
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+from eemkit.core.history import get_history, has_step
+from eemkit.errors import EEMKitWarning, IndexComputationError
+
+LN10 = float(np.log(10.0))
+
+#: Default slope ranges in nm (Helms et al. 2008).
+SLOPE_RANGES: dict[str, tuple[float, float]] = {"S275_295": (275.0, 295.0),
+                                                "S350_400": (350.0, 400.0)}
+
+#: Default wavelengths in nm.
+COEFFICIENTS: tuple[float, ...] = (254.0, 350.0)
+E2_E3: tuple[float, float] = (250.0, 365.0)
+
+COLUMNS = ("a254", "a350", "S275_295", "S350_400", "SR", "E2_E3", "SUVA254")
+
+
+def baseline_corrected(absorbance: xr.DataArray,
+                       baseline: tuple[float, float] | None) -> xr.DataArray:
+    """Each spectrum minus its mean over ``baseline`` (nm); unchanged if ``None``."""
+    if baseline is None:
+        return absorbance
+    lo, hi = baseline
+    wl = absorbance["wavelength"].values
+    inside = (wl >= lo) & (wl <= hi)
+    if not inside.any():
+        raise IndexComputationError(f"baseline range {lo}-{hi} nm is outside the measured "
+                                    f"{wl.min()}-{wl.max()} nm")
+    offset = absorbance.isel(wavelength=np.flatnonzero(inside)).mean("wavelength")
+    return absorbance - offset
+
+
+def _at(spec: np.ndarray, wl: np.ndarray, lam: float) -> float:
+    """Linear interpolation of one spectrum at ``lam`` nm (no extrapolation)."""
+    if lam < wl.min() or lam > wl.max():
+        raise IndexComputationError(f"{lam} nm is outside the measured {wl.min()}-{wl.max()} nm")
+    order = np.argsort(wl)
+    return float(np.interp(lam, wl[order], spec[order]))
+
+
+def spectral_slope(a: np.ndarray, wl: np.ndarray, lo: float, hi: float,
+                   min_r2: float) -> tuple[float, float, str]:
+    """Slope S (nm^-1, positive for decreasing a) of ln a against wavelength on [lo, hi].
+
+    Linear least squares on the measured points in the range (Helms et al.
+    2008). Returns ``(S, R2, note)``; ``S`` is NaN with a note when a value in
+    the range is not positive, when there are fewer than 3 points, or when
+    ``R2 < min_r2``.
+    """
+    sel = (wl >= lo) & (wl <= hi)
+    x, y = wl[sel], a[sel]
+    name = f"{lo:g}-{hi:g} nm"
+    if x.size < 3:
+        return np.nan, np.nan, f"fewer than 3 points in {name}"
+    if not np.all(np.isfinite(y)):
+        return np.nan, np.nan, f"missing absorbance in {name}"
+    if np.any(y <= 0):
+        return np.nan, np.nan, f"a <= 0 in {name}"
+    ly = np.log(y)
+    slope, icept = np.polyfit(x, ly, 1)
+    resid = ly - (slope * x + icept)
+    ss_tot = float(np.sum((ly - ly.mean()) ** 2))
+    r2 = 1.0 - float(np.sum(resid ** 2)) / ss_tot if ss_tot > 0 else np.nan
+    if not np.isfinite(r2) or r2 < min_r2:
+        return np.nan, r2, f"R2 {r2:.3f} < {min_r2:g} in {name}"
+    return float(-slope), r2, ""
+
+
+def compute_absorbance_metrics(
+    ds: xr.Dataset,
+    *,
+    absorbance_path_length_cm: float,
+    baseline: tuple[float, float] | None = (680.0, 700.0),
+    min_r2: float = 0.95,
+    coefficients: Sequence[float] = COEFFICIENTS,
+    slopes: dict[str, tuple[float, float]] | None = None,
+    e2_e3: tuple[float, float] = E2_E3,
+    doc: str | None = None,
+) -> pd.DataFrame:
+    """CDOM metrics, one row per sample, with the settings in ``df.attrs``.
+
+    Parameters
+    ----------
+    ds
+        Dataset with ``absorbance(sample, wavelength)``, blank-corrected.
+    absorbance_path_length_cm
+        Path length of the absorbance cuvette. Required: eemkit never assumes
+        a cuvette.
+    baseline
+        Wavelength range (nm) whose mean is subtracted from each spectrum;
+        ``None`` for no baseline correction.
+    min_r2
+        Smallest R^2 of the log-linear fit for a slope to be reported. This is
+        an eemkit quality threshold, not a value from the literature; report
+        the value used.
+    coefficients
+        Wavelengths (nm) of the Napierian absorption coefficients ``a<nm>``.
+    slopes
+        ``{name: (lo, hi)}``; default :data:`SLOPE_RANGES`. ``SR`` is computed
+        when both default ranges are present.
+    e2_e3
+        The two wavelengths of the E2:E3 ratio.
+    doc
+        Name of a ``sample`` coordinate holding DOC in mg L^-1; ``SUVA254``
+        is computed only when it is given.
+
+    Notes
+    -----
+    Values that cannot be computed are NaN and the reasons are joined in the
+    ``cdom_notes`` column, so low-signal spectra are never reported as
+    numbers.
+    """
+    if "absorbance" not in ds:
+        raise IndexComputationError("no absorbance in the dataset")
+    if not has_step(ds, "subtract_absorbance_blank"):
+        raise IndexComputationError("subtract the absorbance blank first "
+                                    "(eemkit.corrections.subtract_absorbance_blank)")
+    if not np.isfinite(absorbance_path_length_cm) or absorbance_path_length_cm <= 0:
+        raise IndexComputationError("absorbance_path_length_cm must be a positive number")
+    if doc is not None and doc not in ds.coords:
+        raise IndexComputationError(f"no sample column {doc!r} for DOC")
+    slopes = dict(SLOPE_RANGES if slopes is None else slopes)
+    path_m = absorbance_path_length_cm / 100.0
+
+    A = baseline_corrected(ds["absorbance"].transpose("sample", "wavelength"), baseline)
+    wl = A["wavelength"].values
+    rows, notes = [], []
+    for i in range(A.sizes["sample"]):
+        spec = A.values[i]
+        row: dict[str, float] = {}
+        why: list[str] = []
+        if not np.any(np.isfinite(spec)):
+            notes.append("no absorbance")
+            rows.append({})
+            continue
+        for lam in coefficients:
+            v = _at(spec, wl, lam)
+            ok = v > 0
+            row[f"a{lam:g}"] = LN10 * v / path_m if ok else np.nan
+            if not ok:
+                why.append(f"A({lam:g}) <= 0")
+        a = LN10 * spec / path_m
+        for name, (lo, hi) in slopes.items():
+            s, r2, note = spectral_slope(a, wl, lo, hi, min_r2)
+            row[name], row[f"{name}_R2"] = s, r2
+            if note:
+                why.append(f"{name}: {note}")
+        if {"S275_295", "S350_400"} <= set(slopes):
+            row["SR"] = row["S275_295"] / row["S350_400"]
+        num, den = _at(spec, wl, e2_e3[0]), _at(spec, wl, e2_e3[1])
+        row["E2_E3"] = num / den if num > 0 and den > 0 else np.nan
+        if not (num > 0 and den > 0):
+            why.append(f"E2_E3: A({e2_e3[0]:g}) or A({e2_e3[1]:g}) <= 0")
+        if doc is not None:
+            c = pd.to_numeric(str(ds[doc].values[i]).strip().replace(",", "."),
+                              errors="coerce")
+            a254_dec = _at(spec, wl, 254.0) / path_m
+            ok = np.isfinite(c) and c > 0 and a254_dec > 0
+            row["SUVA254"] = a254_dec / c if ok else np.nan
+            if not ok:
+                why.append("SUVA254: DOC missing or not positive, or A(254) <= 0")
+        rows.append(row)
+        notes.append("; ".join(why))
+
+    df = pd.DataFrame(rows, index=pd.Index(A["sample"].values, name="sample"))
+    df["cdom_notes"] = notes
+    n_bad = int((df["cdom_notes"] != "").sum())
+    if n_bad:
+        warnings.warn(f"CDOM metrics: {n_bad} sample(s) with values left empty (low signal); "
+                      "see the cdom_notes column", EEMKitWarning, stacklevel=2)
+    df.attrs["settings"] = {
+        "absorption_coefficient": "Napierian, a = ln(10) A / l, m^-1 (Green & Blough 1994)",
+        "absorbance_path_length_cm": absorbance_path_length_cm,
+        "baseline_nm": list(baseline) if baseline is not None else None,
+        "slopes": {"ranges_nm": {k: list(v) for k, v in slopes.items()},
+                   "fit": "linear least squares of ln a (Helms et al. 2008)",
+                   "min_r2": min_r2},
+        "E2_E3_nm": list(e2_e3),
+        "SUVA254": {"doc_column": doc, "units": "L mg^-1 m^-1",
+                    "reference": "Weishaar et al. 2003"} if doc else None,
+        "processing_history": get_history(ds),
+    }
+    return df

eemkit/campaigns.py

@@ -0,0 +1,276 @@
+"""Comparing sampling campaigns, sites and rivers.
+
+Works on any per-sample table (indices, CDOM metrics, PARAFAC Fmax) with
+columns that name the *group* (e.g. river), the *site* and the *time* (e.g.
+campaign). The statistical unit is the mean of the replicates of one site at
+one time (:func:`site_means`), so replicates are never counted as independent
+samples.
+
+* :func:`design_table` and :func:`confounding` show which groups were sampled
+  at which times, and say when group and time cannot be separated (for
+  example rivers that were sampled in different months);
+* :func:`compare_times`: within each group, Kruskal-Wallis test between times
+  (Kruskal & Wallis 1952) and Dunn's pairwise test (Dunn 1964) with Holm's
+  correction (Holm 1979);
+* :func:`spatial_trend`: Spearman correlation with a site position (e.g.
+  distance from the source), per group;
+* :func:`covariate_correlations`: Spearman correlation with covariates per
+  group and time (e.g. discharge, rainfall);
+* :func:`composition`: the share of each PARAFAC component in total Fmax.
+
+Optional tables are joined explicitly on their key columns
+(:func:`join_sites`, :func:`join_covariates`); rows that do not match are
+reported, never guessed.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from itertools import combinations
+
+import numpy as np
+import pandas as pd
+from scipy import stats
+
+from eemkit.errors import SettingsError
+
+
+def _check_columns(df: pd.DataFrame, cols: Sequence[str]) -> None:
+    missing = [c for c in cols if c and c not in df.columns]
+    if missing:
+        raise SettingsError(f"no column {missing} in the table (columns: {list(df.columns)})")
+
+
+def _time_order(values) -> list:
+    return sorted(pd.unique(pd.Series(values).astype(str)))
+
+
+# --------------------------------------------------------------------------- #
+# Units and composition
+# --------------------------------------------------------------------------- #
+
+def site_means(df: pd.DataFrame, values: Sequence[str], *, group: str, time: str,
+               site: str | None = None) -> pd.DataFrame:
+    """Mean, standard deviation and coefficient of variation of the replicates.
+
+    One row per (group, site, time); ``n`` is the number of replicates. Without
+    ``site`` all samples of a group at one time are pooled.
+    """
+    keys = [group] + ([site] if site else []) + [time]
+    _check_columns(df, list(keys) + list(values))
+    data = df[keys + list(values)].copy()
+    for k in keys:
+        data[k] = data[k].astype(str)
+    g = data.groupby(keys, sort=True)
+    mean = g[list(values)].mean()
+    sd = g[list(values)].std(ddof=1).add_suffix("_sd")
+    with np.errstate(divide="ignore", invalid="ignore"):
+        cv = (100 * g[list(values)].std(ddof=1) / mean.abs()).add_suffix("_cv_pct")
+    out = pd.concat([mean, sd, cv], axis=1)
+    out.insert(0, "n", g.size())
+    return out.reset_index()
+
+
+def composition(df: pd.DataFrame, fmax_columns: Sequence[str]) -> pd.DataFrame:
+    """Percentage of each component in the summed Fmax of each sample."""
+    _check_columns(df, fmax_columns)
+    f = df[list(fmax_columns)].astype(float)
+    total = f.sum(axis=1, min_count=len(fmax_columns))
+    out = f.div(total, axis=0) * 100
+    out.columns = [f"{c.removesuffix('_Fmax')}_pct" for c in fmax_columns]
+    return out
+
+
+# --------------------------------------------------------------------------- #
+# Design
+# --------------------------------------------------------------------------- #
+
+def design_table(units: pd.DataFrame, *, group: str, time: str) -> pd.DataFrame:
+    """Number of units for every (time, group) combination."""
+    _check_columns(units, [group, time])
+    t = pd.crosstab(units[time].astype(str), units[group].astype(str))
+    return t.reindex(_time_order(units[time]))
+
+
+def confounding(units: pd.DataFrame, *, group: str, time: str) -> list[str]:
+    """Plain-language warnings when group and time effects cannot be separated."""
+    t = design_table(units, group=group, time=time) > 0
+    msgs = []
+    groups = list(t.columns)
+    for a, b in combinations(groups, 2):
+        shared = t.index[t[a] & t[b]].tolist()
+        if not shared:
+            msgs.append(f"{a} and {b} were never sampled at the same {time}: differences "
+                        f"between them cannot be separated from differences between "
+                        f"{time} values.")
+    if len(groups) > 1 and (t.sum(axis=1) <= 1).all():
+        msgs.insert(0, f"Every {time} contains a single {group}: {group} and {time} are "
+                       "completely confounded; compare times within each group only.")
+    return msgs
+
+
+# --------------------------------------------------------------------------- #
+# Tests
+# --------------------------------------------------------------------------- #
+
+def _holm(p: np.ndarray) -> np.ndarray:
+    p = np.asarray(p, dtype=float)
+    out = np.full_like(p, np.nan)
+    ok = np.isfinite(p)
+    if not ok.any():
+        return out
+    pv = p[ok]
+    order = np.argsort(pv)
+    m = pv.size
+    adj = np.minimum(1.0, np.maximum.accumulate((m - np.arange(m)) * pv[order]))
+    res = np.empty(m)
+    res[order] = adj
+    out[ok] = res
+    return out
+
+
+def dunn(values: Sequence[float], labels: Sequence[str]) -> pd.DataFrame:
+    """Dunn's (1964) pairwise test after Kruskal-Wallis, with tie correction and Holm.
+
+    Returns one row per pair: mean ranks, z, two-sided p and Holm-adjusted p.
+    """
+    v = np.asarray(values, dtype=float)
+    lab = np.asarray(labels).astype(str)
+    ok = np.isfinite(v)
+    v, lab = v[ok], lab[ok]
+    n = v.size
+    ranks = stats.rankdata(v)
+    _, counts = np.unique(v, return_counts=True)
+    tie = np.sum(counts ** 3 - counts) / (12.0 * (n - 1)) if n > 1 else 0.0
+    levels = _time_order(lab)
+    rows = []
+    for a, b in combinations(levels, 2):
+        ra, rb = ranks[lab == a], ranks[lab == b]
+        se = np.sqrt((n * (n + 1) / 12.0 - tie) * (1 / ra.size + 1 / rb.size))
+        z = (ra.mean() - rb.mean()) / se if se > 0 else np.nan
+        p = 2 * stats.norm.sf(abs(z)) if np.isfinite(z) else np.nan
+        rows.append({"a": a, "b": b, "n_a": ra.size, "n_b": rb.size, "mean_rank_a": ra.mean(),
+                     "mean_rank_b": rb.mean(), "z": z, "p": p})
+    out = pd.DataFrame(rows)
+    if not out.empty:
+        out["p_holm"] = _holm(out["p"].to_numpy())
+    return out
+
+
+def compare_times(units: pd.DataFrame, value: str, *, group: str,
+                  time: str) -> tuple[pd.DataFrame, pd.DataFrame]:
+    """Within each group, differences of ``value`` between times.
+
+    Returns ``(kruskal, pairwise)``: one Kruskal-Wallis row per group (H,
+    degrees of freedom, p, number of units and times, or a note when the test
+    cannot be run) and Dunn's pairwise comparisons with Holm-adjusted p.
+    """
+    _check_columns(units, [group, time, value])
+    kw_rows, pair_rows = [], []
+    for g, sub in units.groupby(units[group].astype(str), sort=True):
+        sub = sub[np.isfinite(sub[value].astype(float))]
+        samples = [s[value].astype(float).to_numpy()
+                   for _, s in sub.groupby(sub[time].astype(str), sort=True)]
+        row = {group: g, "n_units": len(sub), "n_times": len(samples),
+               "H": np.nan, "df": np.nan, "p": np.nan, "note": ""}
+        if len(samples) < 2:
+            row["note"] = "fewer than 2 times with data"
+        elif len(sub) - len(samples) < 1:
+            row["note"] = "one unit per time: no replication between sites"
+        else:
+            try:
+                h, p = stats.kruskal(*samples)
+                row.update(H=float(h), df=len(samples) - 1, p=float(p))
+            except ValueError as exc:  # e.g. all values identical
+                row["note"] = str(exc)
+        kw_rows.append(row)
+        if np.isfinite(row["p"]):
+            d = dunn(sub[value].astype(float), sub[time].astype(str))
+            d.insert(0, group, g)
+            pair_rows.append(d)
+    kw = pd.DataFrame(kw_rows)
+    pairs = pd.concat(pair_rows, ignore_index=True) if pair_rows else pd.DataFrame()
+    kw.attrs["method"] = ("Kruskal-Wallis between times within each group (Kruskal & Wallis "
+                          "1952); Dunn's pairwise test (Dunn 1964), Holm-adjusted (Holm 1979); "
+                          "unit = mean of the replicates of one site at one time")
+    return kw, pairs
+
+
+def _spearman(x, y) -> dict:
+    x, y = np.asarray(x, dtype=float), np.asarray(y, dtype=float)
+    ok = np.isfinite(x) & np.isfinite(y)
+    n = int(ok.sum())
+    if n < 3 or np.ptp(x[ok]) == 0 or np.ptp(y[ok]) == 0:
+        return {"n": n, "rho": np.nan, "p": np.nan,
+                "note": "fewer than 3 pairs or no variation"}
+    rho, p = stats.spearmanr(x[ok], y[ok])
+    return {"n": n, "rho": float(rho), "p": float(p), "note": ""}
+
+
+def spatial_trend(units: pd.DataFrame, value: str, position: str, *, group: str,
+                  time: str | None = None) -> pd.DataFrame:
+    """Spearman correlation between ``value`` and a site ``position`` within each group.
+
+    With ``time``, also one row per (group, time). ``position`` is any number
+    that orders the sites, e.g. distance from the source in km.
+    """
+    _check_columns(units, [group, value, position] + ([time] if time else []))
+    rows = []
+    for g, sub in units.groupby(units[group].astype(str), sort=True):
+        rows.append({group: g, **({time: "all"} if time else {}),
+                     **_spearman(sub[position], sub[value])})
+        if time:
+            for t, s2 in sub.groupby(sub[time].astype(str), sort=True):
+                rows.append({group: g, time: t, **_spearman(s2[position], s2[value])})
+    return pd.DataFrame(rows)
+
+
+def covariate_correlations(units: pd.DataFrame, value: str, covariates: Sequence[str], *,
+                           group: str) -> pd.DataFrame:
+    """Spearman correlation between ``value`` and each covariate, per group and overall."""
+    _check_columns(units, [group, value] + list(covariates))
+    rows = []
+    for c in covariates:
+        rows.append({group: "all", "covariate": c, **_spearman(units[c], units[value])})
+        for g, sub in units.groupby(units[group].astype(str), sort=True):
+            rows.append({group: g, "covariate": c, **_spearman(sub[c], sub[value])})
+    return pd.DataFrame(rows)
+
+
+# --------------------------------------------------------------------------- #
+# Optional tables
+# --------------------------------------------------------------------------- #
+
+def _join(units: pd.DataFrame, table: pd.DataFrame, keys: Sequence[str],
+          what: str) -> tuple[pd.DataFrame, list[str]]:
+    _check_columns(table, keys)
+    t = table.copy()
+    for k in keys:
+        t[k] = t[k].astype(str).str.strip()
+    dup = t.duplicated(list(keys))
+    if dup.any():
+        raise SettingsError(f"{what}: duplicate rows for {t.loc[dup, list(keys)].values.tolist()}")
+    left = units.copy()
+    for k in keys:
+        left[k] = left[k].astype(str).str.strip()
+    merged = left.merge(t, on=list(keys), how="left", indicator=True)
+    unmatched = merged.loc[merged["_merge"] == "left_only", list(keys)].drop_duplicates()
+    notes = [f"{what}: no row for " + ", ".join(f"{k}={r[k]}" for k in keys)
+             for _, r in unmatched.iterrows()]
+    return merged.drop(columns="_merge"), notes
+
+
+def join_sites(units: pd.DataFrame, sites: pd.DataFrame, *, group: str,
+               site: str) -> tuple[pd.DataFrame, list[str]]:
+    """Add the columns of a site table (e.g. latitude, longitude, distance) by (group, site)."""
+    return _join(units, sites, [group, site], "site table")
+
+
+def join_covariates(units: pd.DataFrame, covariates: pd.DataFrame, *, group: str,
+                    time: str) -> tuple[pd.DataFrame, list[str]]:
+    """Add the columns of a covariate table (e.g. discharge, rainfall) by (group, time).
+
+    The time values must be written exactly as in the sample table (e.g.
+    ``2024-05`` for a campaign); no date matching is attempted.
+    """
+    return _join(units, covariates, [group, time], "covariate table")

eemkit/core/__init__.py

@@ -0,0 +1 @@
+"""Data model: grids, multi-sample datasets and processing history."""

eemkit/core/dataset.py

@@ -0,0 +1,91 @@
+"""Combining single measurements into a multi-sample dataset.
+
+A single measurement (one EEM, one absorbance spectrum) is an
+``xarray.Dataset`` whose attributes describe the file and the instrument.
+:func:`stack` puts many of them along a new ``sample`` dimension:
+
+* data variables gain the ``sample`` dimension; a missing measurement
+  (``None``) becomes all-NaN;
+* coordinates such as ``scan_time(ex)`` become ``(sample, ex)``;
+* attributes equal for every sample stay dataset attributes, the others
+  become ``sample`` coordinates, so nothing is lost;
+* the processing history is kept once if it is the same for every sample.
+
+A ``prefix`` keeps EEM and absorbance metadata apart in the same dataset
+(both have e.g. an ``instrument`` attribute).
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+
+import numpy as np
+import xarray as xr
+
+from eemkit.core.history import HISTORY_ATTR, get_history, same_history
+from eemkit.errors import GridError
+
+SAMPLE_DIM = "sample"
+
+
+def stack(
+    items: Sequence[xr.Dataset | None],
+    sample_ids: Sequence[str],
+    prefix: str = "",
+) -> xr.Dataset:
+    """Stack single measurements that share the same grid along ``sample``."""
+    if len(items) != len(sample_ids):
+        raise ValueError("items and sample_ids differ in length")
+    if len(set(sample_ids)) != len(sample_ids):
+        raise ValueError("sample_ids are not unique")
+    present = [it for it in items if it is not None]
+    if not present:
+        raise ValueError("nothing to stack: every item is None")
+    ref = present[0]
+
+    for it in present[1:]:
+        for dim in ref.dims:
+            if dim not in it.dims or it.sizes[dim] != ref.sizes[dim] or \
+                    not np.allclose(it[dim].values, ref[dim].values):
+                raise GridError(
+                    f"'{it.attrs.get('source_file', '?')}' is not on the same '{dim}' grid as "
+                    f"'{ref.attrs.get('source_file', '?')}'; resample to a common grid first"
+                )
+
+    # index values only: ref[d] would also drag along coordinates such as scan_time(ex)
+    out = xr.Dataset(coords={d: (d, ref[d].values, ref[d].attrs) for d in ref.dims})
+    out = out.assign_coords({SAMPLE_DIM: (SAMPLE_DIM, np.asarray(sample_ids, dtype=object))})
+
+    def stacked(name: str, template: xr.DataArray) -> np.ndarray:
+        fill = np.array("NaT", dtype=template.dtype) if template.dtype.kind == "M" else np.nan
+        return np.stack([it[name].values if it is not None
+                         else np.full(template.shape, fill, dtype=template.dtype)
+                         for it in items])
+
+    for name, var in ref.data_vars.items():
+        out[name] = ((SAMPLE_DIM, *var.dims), stacked(name, var), var.attrs)
+    for name, coord in ref.coords.items():
+        if name in ref.dims:
+            continue
+        out = out.assign_coords(
+            {prefix + name: ((SAMPLE_DIM, *coord.dims), stacked(name, coord), coord.attrs)}
+        )
+
+    keys = [k for k in dict.fromkeys(k for it in present for k in it.attrs) if k != HISTORY_ATTR]
+    for key in keys:
+        vals = [str(it.attrs.get(key, "")) if it is not None else "" for it in items]
+        present_vals = {str(it.attrs.get(key, "")) for it in present}
+        if len(present_vals) == 1 and all(it.attrs.get(key) is not None for it in present):
+            out.attrs[prefix + key] = present[0].attrs[key]
+        else:
+            out = out.assign_coords({prefix + key: (SAMPLE_DIM, np.asarray(vals, dtype=object))})
+
+    histories = [get_history(it) for it in present]
+    if all(same_history(histories[0], h) for h in histories[1:]):
+        out.attrs[HISTORY_ATTR] = json.dumps(histories[0])
+    else:
+        out = out.assign_coords({prefix + "history": (SAMPLE_DIM, np.asarray(
+            [json.dumps(get_history(it)) if it is not None else "" for it in items],
+            dtype=object))})
+    return out