"""Leakage-safe Linear/Ridge baseline with reusable NetCDF model artifacts.""" from __future__ import annotations import argparse import logging import os from datetime import datetime, timezone from pathlib import Path import numpy as np import xarray as xr LOGGER = logging.getLogger(__name__) def _matrix(data: xr.DataArray, sample_dim: str, value_dim: str, label: str) -> xr.DataArray: if set(data.dims) != {sample_dim, value_dim} or data.ndim != 2: raise ValueError(f"{label} must have exactly dimensions {sample_dim!r}, {value_dim!r}") return data.transpose(sample_dim, value_dim).astype(float) def fit_ridge(predictors: xr.DataArray, targets: xr.DataArray, alpha: float, training_start: str, training_end: str, sample_dim: str = "time", feature_dim: str = "feature", target_dim: str = "target") -> xr.Dataset: """Fit standardized Linear (alpha=0) or Ridge on one temporal training slice.""" if not np.isfinite(alpha) or alpha < 0: raise ValueError("alpha must be finite and non-negative") x = _matrix(predictors, sample_dim, feature_dim, "predictors") y = _matrix(targets, sample_dim, target_dim, "targets") x, y = xr.align(x, y, join="exact") if sample_dim not in x.coords or not np.issubdtype(x[sample_dim].dtype, np.datetime64): raise ValueError("Sample dimension must have datetime64 coordinates") x_train = x.sel({sample_dim: slice(training_start, training_end)}) y_train = y.sel({sample_dim: slice(training_start, training_end)}) if x_train.sizes[sample_dim] < 2: raise ValueError("Training slice must contain at least two samples") if not np.isfinite(x_train.values).all() or not np.isfinite(y_train.values).all(): raise ValueError("Training predictors/targets contain missing or non-finite values") feature_mean = x_train.mean(sample_dim) raw_std = x_train.std(sample_dim) constant = raw_std == 0 feature_scale = raw_std.where(~constant, 1.0) x_standard = ((x_train - feature_mean) / feature_scale).values target_mean = y_train.mean(sample_dim) y_centered = y_train.values - target_mean.values if alpha == 0: coefficient, _, _, _ = np.linalg.lstsq(x_standard, y_centered, rcond=None) else: gram = x_standard.T @ x_standard + alpha * np.eye(x_standard.shape[1]) coefficient = np.linalg.solve(gram, x_standard.T @ y_centered) coefficient[constant.values, :] = 0.0 model = xr.Dataset( { "coefficient": ((feature_dim, target_dim), coefficient), "intercept": ((target_dim,), target_mean.values), "feature_mean": ((feature_dim,), feature_mean.values), "feature_scale": ((feature_dim,), feature_scale.values), "constant_feature": ((feature_dim,), constant.values.astype(np.int8)), }, coords={feature_dim: x_train[feature_dim], target_dim: y_train[target_dim]}, ) model.attrs = { "model": "LinearRegression" if alpha == 0 else "Ridge", "alpha": float(alpha), "training_start": str(x_train[sample_dim].min().values), "training_end": str(x_train[sample_dim].max().values), "training_samples": int(x_train.sizes[sample_dim]), "scaler_fit_scope": "training slice only", "missing_policy": "reject non-finite training values", "software": "dr-gcm 0.1.0", "created_utc": datetime.now(timezone.utc).isoformat(), } return model def predict_ridge(predictors: xr.DataArray, model: xr.Dataset, sample_dim: str = "time", feature_dim: str = "feature", target_dim: str = "target") -> xr.DataArray: """Apply a saved Linear/Ridge artifact without refitting preprocessing.""" required = {"coefficient", "intercept", "feature_mean", "feature_scale"} missing = sorted(required - set(model.data_vars)) if missing: raise KeyError(f"Model artifact is missing: {missing}") x = _matrix(predictors, sample_dim, feature_dim, "predictors") if not np.array_equal(x[feature_dim].values, model[feature_dim].values): raise ValueError("Prediction feature coordinate/order differs from fitted model") if not np.isfinite(x.values).all(): raise ValueError("Prediction predictors contain missing or non-finite values") standardized = (x.values - model.feature_mean.values) / model.feature_scale.values values = standardized @ model.coefficient.values + model.intercept.values result = xr.DataArray( values, dims=(sample_dim, target_dim), coords={sample_dim: x[sample_dim], target_dim: model[target_dim]}, name="prediction", attrs={"model": model.attrs.get("model", ""), "model_alpha": model.attrs.get("alpha", np.nan), "model_training_start": model.attrs.get("training_start", ""), "model_training_end": model.attrs.get("training_end", "")}, ) return result def fit_monthly_climatology(observed: xr.DataArray, training_start: str, training_end: str, time_dim: str = "time") -> xr.Dataset: """Fit a 12-month climatology using only the explicit training slice.""" if time_dim not in observed.dims or time_dim not in observed.coords: raise ValueError(f"Observed data require coordinate {time_dim!r}") if not np.issubdtype(observed[time_dim].dtype, np.datetime64): raise ValueError("Climatology time coordinate must be datetime64") training = observed.sel({time_dim: slice(training_start, training_end)}) if training.sizes.get(time_dim, 0) == 0: raise ValueError("Climatology training slice selects no data") if not np.isfinite(training.values).all(): raise ValueError("Climatology training data contain non-finite values") available_months = np.unique(training[time_dim].dt.month.values) if not np.array_equal(available_months, np.arange(1, 13)): raise ValueError("Climatology training slice must contain all 12 calendar months") climatology = training.groupby(f"{time_dim}.month").mean(time_dim).rename("climatology") climatology.attrs = dict(observed.attrs) result = climatology.to_dataset() result.attrs = { "model": "monthly_climatology", "training_start": str(training[time_dim].min().values), "training_end": str(training[time_dim].max().values), "training_samples": int(training.sizes[time_dim]), "fit_scope": "training slice only", "software": "dr-gcm 0.1.0", "created_utc": datetime.now(timezone.utc).isoformat(), } return result def predict_monthly_climatology(valid_time: xr.DataArray, model: xr.Dataset) -> xr.DataArray: """Select saved calendar-month climatology for arbitrary forecast valid times.""" if "climatology" not in model or "month" not in model.climatology.dims: raise KeyError("Climatology model must contain climatology(month, ...)") if valid_time.ndim != 1 or not np.issubdtype(valid_time.dtype, np.datetime64): raise ValueError("valid_time must be a one-dimensional datetime64 coordinate") result = model.climatology.sel(month=valid_time.dt.month) if "month" in result.coords: result = result.drop_vars("month") dimension = valid_time.dims[0] result = result.assign_coords(valid_time=(dimension, valid_time.values)).rename("climatology_prediction") result.attrs = dict(model.climatology.attrs) result.attrs.update({"model": "monthly_climatology", "model_training_start": model.attrs.get("training_start", ""), "model_training_end": model.attrs.get("training_end", "")}) return result def persistence_forecast(observed: xr.DataArray, init: str, valid_time: xr.DataArray, time_dim: str = "time") -> xr.DataArray: """Repeat the observation from the month immediately before initialization.""" if time_dim not in observed.dims or not np.issubdtype(observed[time_dim].dtype, np.datetime64): raise ValueError("Persistence observations require datetime64 time") if valid_time.ndim != 1 or not np.issubdtype(valid_time.dtype, np.datetime64): raise ValueError("Persistence valid_time must be one-dimensional datetime64") init_month = np.datetime64(f"{init}-01", "M") source_month = init_month - np.timedelta64(1, "M") observed_months = observed[time_dim].values.astype("datetime64[M]") matches = np.flatnonzero(observed_months == source_month) if matches.size != 1: raise ValueError(f"Persistence requires exactly one observation for {source_month}") forecast_months = valid_time.values.astype("datetime64[M]") if np.any(forecast_months < init_month): raise ValueError("Persistence target valid_time must not precede initial month") dimension = valid_time.dims[0] source = observed.isel({time_dim: int(matches[0])}, drop=True) result = source.expand_dims({dimension: valid_time[dimension]}).assign_coords( valid_time=(dimension, valid_time.values)).rename("persistence_prediction") result.attrs = dict(observed.attrs) result.attrs.update({"model": "persistence", "source_valid_time": str(source_month.astype("datetime64[ns]")), "initial_time": str(init_month.astype("datetime64[ns]"))}) return result def validate_model(path: Path) -> dict: required = {"coefficient", "intercept", "feature_mean", "feature_scale", "constant_feature"} with xr.open_dataset(path) as dataset: missing = sorted(required - set(dataset.data_vars)) if missing: raise ValueError(f"Baseline model is missing: {missing}") if dataset.attrs.get("scaler_fit_scope") != "training slice only": raise ValueError("Baseline model lacks training-only scaler provenance") return {"variables": sorted(required), "sizes": dict(dataset.sizes), "bytes": path.stat().st_size} def validate_prediction(path: Path, variable: str) -> dict: with xr.open_dataset(path) as dataset: if variable not in dataset: raise ValueError(f"Prediction file is missing {variable!r}") return {"variable": variable, "sizes": dict(dataset[variable].sizes), "bytes": path.stat().st_size} def _atomic(dataset: xr.Dataset, output: Path, validator, overwrite: bool = False) -> tuple[Path, dict]: if output.exists() and not overwrite: return output, validator(output) partial = output.with_suffix(output.suffix + ".part") if partial.exists() and not overwrite: raise FileExistsError(f"Partial output retained: {partial}") if overwrite: output.unlink(missing_ok=True) partial.unlink(missing_ok=True) output.parent.mkdir(parents=True, exist_ok=True) encoding = {name: {"zlib": True, "complevel": 1, "shuffle": True} for name in dataset.data_vars} dataset.to_netcdf(partial, encoding=encoding) report = validator(partial) os.replace(partial, output) return output, report def main() -> None: parser = argparse.ArgumentParser(description=__doc__) commands = parser.add_subparsers(dest="command", required=True) fit = commands.add_parser("fit") predict = commands.add_parser("predict") fit.add_argument("--input", required=True, type=Path) fit.add_argument("--predictor-variable", default="predictors") fit.add_argument("--target-variable", default="targets") fit.add_argument("--alpha", required=True, type=float) fit.add_argument("--training-start", required=True) fit.add_argument("--training-end", required=True) fit.add_argument("--output", required=True, type=Path) predict.add_argument("--input", required=True, type=Path) predict.add_argument("--predictor-variable", default="predictors") predict.add_argument("--model", required=True, type=Path) predict.add_argument("--output-variable", default="prediction") predict.add_argument("--output", required=True, type=Path) for child in (fit, predict): child.add_argument("--sample-dimension", default="time") child.add_argument("--feature-dimension", default="feature") child.add_argument("--target-dimension", default="target") child.add_argument("--overwrite", action="store_true") args = parser.parse_args() logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") with xr.open_dataset(args.input) as source: if args.predictor_variable not in source: raise KeyError(f"Predictor variable {args.predictor_variable!r} not found") predictors = source[args.predictor_variable].load() if args.command == "fit": if args.target_variable not in source: raise KeyError(f"Target variable {args.target_variable!r} not found") targets = source[args.target_variable].load() if args.command == "fit": model = fit_ridge(predictors, targets, args.alpha, args.training_start, args.training_end, args.sample_dimension, args.feature_dimension, args.target_dimension) target, report = _atomic(model, args.output, validate_model, args.overwrite) else: with xr.open_dataset(args.model) as source: model = source.load() prediction = predict_ridge(predictors, model, args.sample_dimension, args.feature_dimension, args.target_dimension).rename(args.output_variable) target, report = _atomic( prediction.to_dataset(), args.output, lambda path: validate_prediction(path, args.output_variable), args.overwrite, ) LOGGER.info("Validated %s: %s", target, report) if __name__ == "__main__": main()