"""Data loading for the trust study, built on xpectra.pipeline."""
from __future__ import annotations
from pathlib import Path
from typing import Any, Iterator
import numpy as np
import pandas as pd
from .pipeline.config import ORDERED_GROUPS, ROUTE_FILES, PreprocessConfig
from .cache import cache_key, memoize
from .config import TrustConfig
def _dedupe_keep_mask(X: np.ndarray, study: np.ndarray) -> np.ndarray:
"""Keep-first boolean mask dropping exact-duplicate spectra within each study.
Rows are rounded before hashing so pure float noise is not treated as a real
difference. Duplicates are keyed by (study, rounded spectrum), so an identical
spectrum appearing in two different studies is never collapsed.
"""
rounded = pd.DataFrame(np.round(X, 6))
rounded.insert(0, "__study__", study)
return ~rounded.duplicated(keep="first").to_numpy()
[docs]
def load_labeled(
route: str,
trust: TrustConfig,
pre: PreprocessConfig | None = None,
) -> dict[str, Any]:
"""Labeled spectra (4 studies, unknowns dropped), grouped 6-class labels.
When ``trust.dedupe_spectra`` is set, exact-duplicate spectra within a study
are dropped (keep-first) before the arrays are returned; ``n_dropped_duplicates``
records how many rows were removed.
"""
pre = pre or PreprocessConfig()
def build() -> dict[str, Any]:
from .pipeline.core import load_training_dataframe, training_arrays
df = load_training_dataframe(route, trust.processed_dir, pre)
arrays = training_arrays(df, pre)
arrays["study"] = df["study"].astype(str).to_numpy()
arrays["environmental"] = (df["environmental"].astype(str) == "Y").to_numpy()
arrays["sample_id"] = df[pre.sample_id_column].astype(str).to_numpy()
arrays["type_original"] = df["type_original"].astype(str).to_numpy()
n_rows = arrays["X_raw"].shape[0]
dropped = 0
if trust.dedupe_spectra:
keep = _dedupe_keep_mask(arrays["X_raw"], arrays["study"])
dropped = int((~keep).sum())
if dropped:
for name, value in list(arrays.items()):
if isinstance(value, np.ndarray) and value.shape[:1] == (n_rows,):
arrays[name] = value[keep]
arrays["n_dropped_duplicates"] = dropped
return arrays
key = cache_key("labeled", route, dedupe=trust.dedupe_spectra)
return memoize(key, build, trust.cache_dir)
[docs]
def load_unknown(
route: str,
feature_columns: list[str],
wavenumbers: np.ndarray,
trust: TrustConfig,
pre: PreprocessConfig | None = None,
) -> dict[str, Any]:
"""The kedzierski_2019_u unlabeled environmental pool, on the labeled grid."""
pre = pre or PreprocessConfig()
def build() -> dict[str, Any]:
from .pipeline.core import metadata_frame, read_csv, spectral_columns_sorted
df = read_csv(Path(trust.processed_dir) / ROUTE_FILES[route])
df = df[df["study"] == pre.unknown_study].copy()
if df.empty:
raise ValueError(f"No rows for unknown study '{pre.unknown_study}' in route '{route}'.")
own_columns, own_wn = spectral_columns_sorted(df, pre.wn_min, pre.wn_max)
if own_columns != list(feature_columns) or not np.array_equal(own_wn, wavenumbers):
raise AssertionError(
"Unknown-pool wavenumber grid does not match the labeled training grid."
)
X = df[feature_columns].to_numpy(dtype=float)
X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0)
return {
"X_raw": X,
"meta": metadata_frame(df),
"sample_id": df[pre.sample_id_column].astype(str).to_numpy(),
}
key = cache_key("unknown", route)
return memoize(key, build, trust.cache_dir)
[docs]
def load_external(
study_id: str,
feature_columns: list[str],
label_encoder: Any,
trust: TrustConfig,
pre: PreprocessConfig | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""One external augmentation pool (openspecy / simple / awi2024) from the norm route.
The combined route files carry all datasets; externals are selected by their ``study``
id and are never part of ``load_labeled``. Returns (X, y, source) on the labeled
feature grid, with labels already in the 6-group scheme encoded by ``label_encoder``.
``source`` is the per-row source_study (finer than the dataset id only for openspecy),
joined from the catalog DB; it falls back to the dataset id if the catalog is absent.
"""
pre = pre or PreprocessConfig()
if study_id not in pre.external_studies:
raise ValueError(f"Unknown external study '{study_id}'. "
f"Expected one of {pre.external_studies}.")
def build() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
from .pipeline.core import read_csv
df = read_csv(Path(trust.processed_dir) / ROUTE_FILES["norm"])
df = df[df["study"] == study_id]
if df.empty:
raise ValueError(
f"No rows for external study '{study_id}' in the norm route. "
"Rebuild the processed data with the step-0 notebook or "
"scripts/process_raw_data.py (externals are included by default).")
X = np.nan_to_num(df[feature_columns].to_numpy(float))
y = label_encoder.transform(df["type"].astype(str).to_numpy())
from .pipeline.build_db import CATALOG_PATH
if Path(CATALOG_PATH).exists():
import sqlite3
con = sqlite3.connect(CATALOG_PATH)
try:
mapping = dict(con.execute(
"SELECT sample_id, source_study FROM spectra WHERE dataset_id=?",
(study_id,)).fetchall())
finally:
con.close()
source = df["sample_id"].astype(str).map(mapping).fillna(study_id).to_numpy()
else:
source = np.full(len(df), study_id)
return X, y, source
key = cache_key("external", study_id, features=tuple(feature_columns),
classes=tuple(label_encoder.classes_))
return memoize(key, build, trust.cache_dir)
[docs]
def loso_splits(study: np.ndarray) -> Iterator[tuple[str, np.ndarray, np.ndarray]]:
"""Leave-one-study-out folds over the labeled studies."""
for held in sorted(np.unique(study)):
test_mask = study == held
yield held, np.flatnonzero(~test_mask), np.flatnonzero(test_mask)
[docs]
def leave_one_class_out(y: np.ndarray, held_class: int) -> tuple[np.ndarray, np.ndarray]:
"""Masks for synthetic OOD: train without one class, hold it out as OOD."""
ood_mask = y == held_class
return ~ood_mask, ood_mask
__all__ = [
"ORDERED_GROUPS",
"load_labeled",
"load_unknown",
"load_external",
"loso_splits",
"leave_one_class_out",
]