Spaces:
Running
Running
| """Read individual examples and align forecasts by their full sample key.""" | |
| from __future__ import annotations | |
| import json | |
| import math | |
| from pathlib import Path | |
| from typing import Sequence | |
| import pyarrow.parquet as pq | |
| from src.config import SAMPLE_KEYS, Method, load_methods | |
| def keyed_rows(table) -> dict[tuple, int]: | |
| columns = [table[name].to_pylist() for name in SAMPLE_KEYS] | |
| keys = list(zip(*columns)) | |
| if any(any(value is None for value in key) for key in keys): | |
| raise ValueError("Null sample keys") | |
| result = {key: index for index, key in enumerate(keys)} | |
| if len(result) != len(keys): | |
| raise ValueError("Duplicate sample keys") | |
| return result | |
| def clean_series(values: list) -> list: | |
| # Keep missing observations at their original time positions. | |
| return [float(value) if value is not None and math.isfinite(value) else None for value in values] | |
| class SampleRepository: | |
| def __init__(self, directory: Path, methods: Sequence[Method] | None = None): | |
| self.directory = directory | |
| self.methods = load_methods(directory) if methods is None else tuple(methods) | |
| self._evaluations = {} | |
| self.samples = pq.read_table(directory / "data" / "structured_effects.parquet") | |
| self.index = keyed_rows(self.samples) | |
| if not self.index: | |
| raise ValueError("No examples available") | |
| micro = self.samples["micro"].to_pylist() | |
| self.entries = [ | |
| {"id": row, **dict(zip(SAMPLE_KEYS, key)), "effects": sorted({ | |
| "/".join(str(event[field]) for field in ("family", "subtype") if event.get(field)) | |
| for event in json.loads(micro[row]) | |
| })} | |
| for key, row in self.index.items() | |
| ] | |
| def catalog(self) -> dict: | |
| return {"samples": self.entries, "methods": [method.to_dict() for method in self.methods]} | |
| def _evaluation(self, method: Method): | |
| if method.id not in self._evaluations: | |
| table = pq.read_table(self.directory / method.relative_path, | |
| columns=SAMPLE_KEYS + ["forecast", "mase"]) | |
| index = keyed_rows(table) | |
| if index.keys() != self.index.keys(): | |
| raise ValueError(f"{method.filename}: sample keys do not match the examples") | |
| self._evaluations[method.id] = table, index | |
| return self._evaluations[method.id] | |
| def sample(self, sample_id: int) -> dict: | |
| if sample_id < 0 or sample_id >= len(self.samples): | |
| raise KeyError(sample_id) | |
| result = {name: self.samples[name][sample_id].as_py() for name in self.samples.column_names} | |
| key = tuple(result[name] for name in SAMPLE_KEYS) | |
| result["id"] = sample_id | |
| result["macro"] = json.loads(result["macro"]) | |
| result["micro"] = json.loads(result["micro"]) | |
| for name in ("history", "future", "tsfm_average_forecast"): | |
| result[name] = clean_series(result[name]) | |
| if len(result["history"]) != result["n_history"] or len(result["future"]) != result["n_horizon"]: | |
| raise ValueError("Series lengths do not match example metadata") | |
| forecasts = [] | |
| for method in self.methods: | |
| table, index = self._evaluation(method) | |
| row = index[key] | |
| values = clean_series(table["forecast"][row].as_py()) | |
| mase = table["mase"][row].as_py() | |
| if len(values) != result["n_horizon"]: | |
| raise ValueError(f"{method.filename}: forecast length does not match the horizon") | |
| if mase is None or not math.isfinite(mase) or mase < 0: | |
| raise ValueError(f"{method.filename}: invalid MASE") | |
| forecasts.append({**method.to_dict(), "forecast": values, "mase": float(mase)}) | |
| result["forecasts"] = forecasts | |
| return result | |