"""Create reproducible rolling-origin year folds without temporal leakage.""" from __future__ import annotations import argparse import json import os from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path import numpy as np import xarray as xr @dataclass(frozen=True) class TemporalFold: fold: int train_years: np.ndarray validation_years: np.ndarray test_years: np.ndarray def masks(self, sample_years: np.ndarray) -> dict[str, np.ndarray]: years = np.asarray(sample_years, dtype=int) return { "train": np.isin(years, self.train_years), "validation": np.isin(years, self.validation_years), "test": np.isin(years, self.test_years), } def as_dict(self, sample_years: np.ndarray | None = None) -> dict: payload = { "fold": self.fold, "train_years": self.train_years.tolist(), "validation_years": self.validation_years.tolist(), "test_years": self.test_years.tolist(), "train_range": [int(self.train_years[0]), int(self.train_years[-1])], "validation_range": [int(self.validation_years[0]), int(self.validation_years[-1])], "test_range": [int(self.test_years[0]), int(self.test_years[-1])], } if sample_years is not None: payload["sample_counts"] = { name: int(mask.sum()) for name, mask in self.masks(sample_years).items() } return payload def rolling_origin_folds(sample_years: np.ndarray, min_train_years: int, validation_years: int, test_years: int, step_years: int, gap_years: int = 0) -> list[TemporalFold]: """Return expanding-train folds with later validation and test blocks.""" values = np.asarray(sample_years) if values.ndim != 1 or values.size == 0: raise ValueError("sample_years must be a non-empty one-dimensional array") if not np.all(np.isfinite(values)) or not np.allclose(values, np.round(values)): raise ValueError("sample_years must contain finite integer years") years = np.unique(values.astype(int)) if np.any(np.diff(years) != 1): raise ValueError("Available years must be contiguous; missing years require explicit resolution") settings = { "min_train_years": min_train_years, "validation_years": validation_years, "test_years": test_years, "step_years": step_years, } if any(not isinstance(value, (int, np.integer)) or value < 1 for value in settings.values()): raise ValueError(f"Fold lengths must be positive integers: {settings}") if not isinstance(gap_years, (int, np.integer)) or gap_years < 0: raise ValueError("gap_years must be a non-negative integer") folds = [] train_count = int(min_train_years) fold_number = 0 while True: validation_start = train_count + gap_years validation_end = validation_start + validation_years test_start = validation_end + gap_years test_end = test_start + test_years if test_end > years.size: break fold = TemporalFold( fold=fold_number, train_years=years[:train_count].copy(), validation_years=years[validation_start:validation_end].copy(), test_years=years[test_start:test_end].copy(), ) if not (fold.train_years[-1] < fold.validation_years[0] and fold.validation_years[-1] < fold.test_years[0]): raise AssertionError("Temporal fold ordering invariant failed") folds.append(fold) fold_number += 1 train_count += step_years if not folds: needed = min_train_years + validation_years + test_years + 2 * gap_years raise ValueError(f"Not enough years for one fold: available={years.size}, needed={needed}") return folds def years_from_coordinate(coordinate: xr.DataArray) -> np.ndarray: values = np.asarray(coordinate.values) if not np.issubdtype(values.dtype, np.datetime64): raise ValueError("Time coordinate must decode to datetime64") return values.astype("datetime64[Y]").astype(int) + 1970 def build_plan(input_path: Path, time_coordinate: str, min_train_years: int, validation_years: int, test_years: int, step_years: int, gap_years: int = 0) -> dict: with xr.open_dataset(input_path) as dataset: if time_coordinate not in dataset.coords: raise KeyError(f"Coordinate {time_coordinate!r} not found in {input_path}") sample_years = years_from_coordinate(dataset[time_coordinate]) folds = rolling_origin_folds(sample_years, min_train_years, validation_years, test_years, step_years, gap_years) return { "method": "expanding-window rolling-origin by complete calendar years", "source_file": str(input_path), "time_coordinate": time_coordinate, "settings": { "min_train_years": min_train_years, "validation_years": validation_years, "test_years": test_years, "step_years": step_years, "gap_years": gap_years, }, "available_year_range": [int(sample_years.min()), int(sample_years.max())], "folds": [fold.as_dict(sample_years) for fold in folds], "leakage_invariant": "max(train) Path: if output.exists() and not overwrite: existing = json.loads(output.read_text(encoding="utf-8")) if existing.get("method") != plan.get("method") or existing.get("settings") != plan.get("settings"): raise FileExistsError(f"Existing CV plan differs: {output}; use --overwrite") return output output.parent.mkdir(parents=True, exist_ok=True) partial = output.with_suffix(output.suffix + ".part") partial.write_text(json.dumps(plan, indent=2, sort_keys=True) + "\n", encoding="utf-8") os.replace(partial, output) return output def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--input", required=True, type=Path) parser.add_argument("--time-coordinate", default="time") parser.add_argument("--min-train-years", required=True, type=int) parser.add_argument("--validation-years", required=True, type=int) parser.add_argument("--test-years", required=True, type=int) parser.add_argument("--step-years", required=True, type=int) parser.add_argument("--gap-years", type=int, default=0) parser.add_argument("--output", required=True, type=Path) parser.add_argument("--overwrite", action="store_true") args = parser.parse_args() plan = build_plan(args.input, args.time_coordinate, args.min_train_years, args.validation_years, args.test_years, args.step_years, args.gap_years) print(write_plan(plan, args.output, args.overwrite)) if __name__ == "__main__": main()