Skills › Scientific Agent Skills › relsa-severity-assessment › forecast_relsa.py
forecast_relsa.py
793 lines · skills/relsa-severity-assessment/scripts/forecast_relsa.py at 5f2bc03
1#!/usr/bin/env python3
2"""foRcast — ARIMA forecasting of RELSA severity trajectories.
4Independent nonseasonal approximation of the published foRcast methodology
5(Lutscher et al., 2026, Front. Physiol. 17:1869563), using statsmodels SARIMAX.
6Predicts a score at a specified future observation time, not time-to-endpoint.
7Interpolation adds no independent information and may understate uncertainty.
8A forecast cannot override observed welfare signs or approved humane endpoints.
9"""
11from __future__ import annotations
13import argparse
14import sys
15import warnings
16from dataclasses import dataclass, field
17from pathlib import Path
18from typing import Iterable, Sequence
20import numpy as np
21import pandas as pd
23sys.path.insert(0, str(Path(__file__).resolve().parent))
25from _common import ( # noqa: E402
26 ID_COL,
27 TIME_COL,
28 ForecastMetrics,
29 forecast_metrics,
30 parse_list,
31 read_relsa_table,
32 validate,
33)
35DEFAULT_STEP = 0.1
36MIN_TRAIN_POINTS = 4
39def _trajectory(times, values) -> tuple[np.ndarray, np.ndarray]:
40 t, y = np.asarray(list(times), dtype=float), np.asarray(list(values), dtype=float)
41 if t.ndim != 1 or t.shape != y.shape or not np.isfinite(t).all() or np.isinf(y).any():
42 raise ValueError("times/values must be matching vectors; times finite, values finite or NaN")
43 order = np.argsort(t)
44 t, y = t[order], y[order]
45 if np.any(np.diff(t) <= 0):
46 raise ValueError("duplicate observation times")
47 return t, y
50# --------------------------------------------------------------------------- #
51# interpolation
52# --------------------------------------------------------------------------- #
53def interpolate_series(
54 times: Iterable[float],
55 values: Iterable[float],
56 step: float = DEFAULT_STEP,
57) -> tuple[np.ndarray, np.ndarray]:
58 """Linearly interpolate a sparse trajectory onto a regular ``step`` grid.
60 Leading and trailing missing values are dropped rather than extrapolated.
61 Returns the grid times and interpolated values.
62 """
63 t, y = _trajectory(times, values)
64 if not np.isfinite(step) or step <= 0:
65 raise ValueError("interpolation step must be finite and positive")
66 ok = np.isfinite(t) & np.isfinite(y)
67 t, y = t[ok], y[ok]
68 if t.size < 2:
69 raise ValueError("need at least 2 observed points to interpolate")
70 order = np.argsort(t)
71 t, y = t[order], y[order]
72 offsets = (t - t[0]) / step
73 if not np.allclose(offsets, np.round(offsets), rtol=0, atol=1e-7):
74 raise ValueError("observation times must align with the interpolation grid")
75 if offsets[-1] > 100000:
76 raise ValueError("interpolation grid exceeds 100000 steps")
77 grid = t[0] + np.arange(int(round(offsets[-1])) + 1) * step
78 return grid, np.interp(grid, t, y)
81# --------------------------------------------------------------------------- #
82# automatic ARIMA
83# --------------------------------------------------------------------------- #
84@dataclass
85class ArimaFit:
86 """A selected ARIMA model and the information criterion that chose it."""
88 order: tuple[int, int, int]
89 drift: bool
90 aicc: float
91 aic: float
92 n_obs: int
93 results: object = field(repr=False, default=None)
95 def label(self) -> str:
96 p, d, q = self.order
97 term = " with intercept" if d == 0 else " with drift"
98 return f"ARIMA({p},{d},{q})" + (term if self.drift else "")
101def _select_d(y: np.ndarray, max_d: int) -> int:
102 """Choose the differencing order by successive KPSS tests, as auto.arima does."""
103 from statsmodels.tsa.stattools import kpss
105 series = np.asarray(y, dtype=float)
106 for d in range(max_d + 1):
107 if series.size < 8 or np.allclose(series, series[0]):
108 return d
109 try:
110 with warnings.catch_warnings():
111 warnings.simplefilter("ignore")
112 p_value = kpss(series, regression="c", nlags="auto")[1]
113 except Exception:
114 return d
115 if p_value > 0.05: # KPSS null is stationarity -> stop differencing
116 return d
117 series = np.diff(series)
118 return max_d
121def _fit_one(
122 y: np.ndarray, order: tuple[int, int, int], drift: bool
123) -> ArimaFit | None:
124 from statsmodels.tsa.statespace.sarimax import SARIMAX
126 k = sum(order[::2]) + (1 if drift else 0) + 1
127 if y.size - order[1] <= k + 2:
128 return None
129 try:
130 with warnings.catch_warnings():
131 warnings.simplefilter("ignore")
132 model = SARIMAX(
133 y,
134 order=order,
135 trend="c" if drift else "n",
136 enforce_stationarity=True,
137 enforce_invertibility=True,
138 )
139 res = model.fit(disp=False, maxiter=200)
140 except Exception:
141 return None
142 aic, aicc = float(res.aic), float(res.aicc)
143 if not res.mle_retvals.get("converged", False) or not np.isfinite([aic, aicc]).all():
144 return None
145 n = int(res.nobs)
146 return ArimaFit(order=order, drift=drift, aicc=aicc, aic=aic,
147 n_obs=n, results=res)
150def auto_arima(
151 y: Iterable[float],
152 max_p: int = 5,
153 max_q: int = 5,
154 max_d: int = 2,
155 d: int | None = None,
156 allow_drift: bool = True,
157 stepwise: bool = True,
158) -> ArimaFit:
159 """Select an ARIMA model by minimising AICc (Hyndman–Khandakar stepwise).
161 Approximates the nonseasonal search in ``forecast::auto.arima``: KPSS, then
162 four seed models — (2,d,2), (0,d,0), (1,d,0), (0,d,1) — followed by a
163 hill-climb over neighbouring (p, q) and the drift term. ``stepwise=False``
164 searches the full grid, which is slower and rarely changes the answer.
166 The best model found is not guaranteed to be the globally best one: the
167 search covers a bounded range of p, d, q, which the paper notes as a
168 limitation of the approach rather than of one implementation.
169 """
170 series = np.asarray(list(y), dtype=float)
171 if series.ndim != 1 or not np.isfinite(series).all():
172 raise ValueError("ARIMA input must be a finite, regular series")
173 if any(int(v) != v or v < 0 for v in (max_p, max_q, max_d)):
174 raise ValueError("ARIMA maximum orders must be nonnegative integers")
175 if d is not None and (int(d) != d or d < 0 or d > max_d):
176 raise ValueError("d must be an integer between zero and max_d")
177 if series.size < MIN_TRAIN_POINTS:
178 raise ValueError(
179 f"need at least {MIN_TRAIN_POINTS} finite points to fit ARIMA, got {series.size}"
180 )
182 order_d = _select_d(series, max_d) if d is None else int(d)
183 drifts = (True, False) if allow_drift and order_d <= 1 else (False,)
185 tried: dict[tuple[tuple[int, int, int], bool], ArimaFit | None] = {}
187 def attempt(order: tuple[int, int, int], drift: bool) -> ArimaFit | None:
188 if order[0] > max_p or order[2] > max_q or min(order[0], order[2]) < 0:
189 return None
190 key = (order, drift)
191 if key not in tried:
192 tried[key] = _fit_one(series, order, drift)
193 return tried[key]
195 candidates: list[ArimaFit] = []
196 if not stepwise:
197 for p in range(max_p + 1):
198 for q in range(max_q + 1):
199 for drift in drifts:
200 fit = attempt((p, order_d, q), drift)
201 if fit:
202 candidates.append(fit)
203 else:
204 seeds = [(2, order_d, 2), (0, order_d, 0), (1, order_d, 0), (0, order_d, 1)]
205 for seed in seeds:
206 for drift in drifts:
207 fit = attempt(seed, drift)
208 if fit:
209 candidates.append(fit)
210 if candidates:
211 best = min(candidates, key=lambda f: f.aicc)
212 improved = True
213 while improved:
214 improved = False
215 p, dd, q = best.order
216 neighbours = [
217 ((p + 1, dd, q), best.drift),
218 ((p - 1, dd, q), best.drift),
219 ((p, dd, q + 1), best.drift),
220 ((p, dd, q - 1), best.drift),
221 ((p + 1, dd, q + 1), best.drift),
222 ((p - 1, dd, q - 1), best.drift),
223 ((p, dd, q), not best.drift),
224 ]
225 for order, drift in neighbours:
226 if drift and (not allow_drift or order_d > 1):
227 continue
228 fit = attempt(order, drift)
229 if fit and fit.aicc < best.aicc - 1e-8:
230 best, improved = fit, True
231 break
232 candidates.append(best)
234 if not candidates:
235 # Nothing converged (usually a near-constant or very short series):
236 # Try the simplest model at the selected differencing order.
237 fit = _fit_one(series, (0, order_d, 0), False)
238 if fit is None:
239 raise RuntimeError("no ARIMA model could be fitted to this series")
240 return fit
241 return min(candidates, key=lambda f: f.aicc)
244# --------------------------------------------------------------------------- #
245# forecasting one animal
246# --------------------------------------------------------------------------- #
247@dataclass
248class Forecast:
249 """A forecast of one animal's RELSA trajectory."""
251 animal: str
252 fit: ArimaFit
253 times: np.ndarray
254 predicted: np.ndarray
255 lower: np.ndarray
256 upper: np.ndarray
257 actual: np.ndarray | None = None
258 train_times: np.ndarray | None = None
259 train_values: np.ndarray | None = None
260 interpolated: bool = True
261 alpha: float = 0.05
262 warnings: list[str] = field(default_factory=list)
264 def metrics(self) -> ForecastMetrics:
265 if self.actual is None:
266 raise ValueError("no actual values to score this forecast against")
267 return forecast_metrics(self.actual, self.predicted, self.lower, self.upper)
269 def to_frame(self) -> pd.DataFrame:
270 data = {
271 ID_COL: self.animal,
272 TIME_COL: self.times,
273 "predicted": self.predicted,
274 "lower": self.lower,
275 "upper": self.upper,
276 "model": self.fit.label(),
277 "alpha": self.alpha,
278 "interpolated": self.interpolated,
279 "n_observed": len(self.train_times) if self.train_times is not None else None,
280 "notes": "; ".join(self.warnings),
281 }
282 if self.actual is not None:
283 data["actual"] = self.actual
284 return pd.DataFrame(data)
287def forecast_animal(
288 times: Iterable[float],
289 values: Iterable[float],
290 target_times: Iterable[float] | None = None,
291 animal: str = "",
292 interpolate_step: float | None = DEFAULT_STEP,
293 alpha: float = 0.05,
294 clip_at_zero: bool = True,
295 **arima_kwargs,
296) -> Forecast:
297 """Fit ARIMA to one RELSA trajectory and forecast the requested time points.
299 ``target_times`` defaults to one observation-spacing step beyond the last
300 training point. With ``interpolate_step`` set, the model is fitted on the
301 interpolated grid and the horizon is converted to grid steps, so targets
302 must lie exactly on that grid.
304 ``clip_at_zero`` floors the forecast and its interval at 0, since RELSA is
305 non-negative by construction; the raw Gaussian interval can dip below.
306 """
307 t, y = _trajectory(times, values)
308 if not np.isfinite(alpha) or not 0 < alpha < 1:
309 raise ValueError("alpha must be between zero and one")
310 if clip_at_zero and np.any(y < 0):
311 raise ValueError("observed RELSA scores must be nonnegative")
312 if interpolate_step is not None and (not np.isfinite(interpolate_step) or interpolate_step <= 0):
313 raise ValueError("interpolate_step must be positive or None")
314 ok = np.isfinite(t) & np.isfinite(y)
315 t, y = t[ok], y[ok]
316 order = np.argsort(t)
317 t, y = t[order], y[order]
318 notes: list[str] = []
320 if t.size < 2:
321 raise ValueError(f"animal {animal!r}: fewer than 2 observed RELSA scores")
323 spacing = float(np.median(np.diff(t))) if t.size > 1 else 1.0
324 if target_times is None:
325 targets = np.array([t[-1] + spacing], dtype=float)
326 else:
327 targets = np.asarray(list(target_times), dtype=float)
328 if targets.ndim != 1 or not targets.size or not np.isfinite(targets).all():
329 raise ValueError("target times must be a nonempty finite vector")
330 if np.any(targets <= t[-1]):
331 raise ValueError(
332 f"animal {animal!r}: target times must be after the last training point "
333 f"({t[-1]}); got {targets.tolist()}"
334 )
336 if interpolate_step:
337 grid, series = interpolate_series(t, y, step=interpolate_step)
338 step = interpolate_step
339 origin = grid[-1]
340 notes.append(
341 f"fitted on {series.size} points interpolated at {interpolate_step} "
342 f"time units from {t.size} observations; prediction intervals are "
343 "conditional on the interpolated series, not independent observations"
344 )
345 else:
346 if not np.allclose(np.diff(t), spacing, rtol=1e-6, atol=1e-9):
347 raise ValueError("uneven observation times require an explicit regular interpolation grid")
348 series = y
349 step = spacing
350 origin = t[-1]
352 if t.size < MIN_TRAIN_POINTS:
353 notes.append(
354 f"only {t.size} observed points before the forecast — the paper's "
355 "own weakest predictions come from exactly this situation; treat the "
356 "interval, not the point estimate, as the message"
357 )
359 offsets = (targets - origin) / step
360 if not np.allclose(offsets, np.round(offsets), rtol=0, atol=1e-7):
361 raise ValueError("target times must align exactly with forecast grid steps")
362 steps = int(round(offsets.max()))
363 if steps > 100000:
364 raise ValueError("forecast horizon exceeds 100000 steps")
365 fit = auto_arima(series, **arima_kwargs)
366 forecast = fit.results.get_forecast(steps=steps) # type: ignore[union-attr]
367 mean = np.asarray(forecast.predicted_mean, dtype=float)
368 conf = np.asarray(forecast.conf_int(alpha=alpha), dtype=float)
370 idx = np.round(offsets).astype(int) - 1
371 predicted, lower, upper = mean[idx], conf[idx, 0], conf[idx, 1]
372 if clip_at_zero:
373 predicted = np.clip(predicted, 0.0, None)
374 lower = np.clip(lower, 0.0, None)
375 upper = np.clip(upper, 0.0, None)
377 return Forecast(
378 animal=animal,
379 fit=fit,
380 times=targets,
381 predicted=predicted,
382 lower=lower,
383 upper=upper,
384 train_times=t,
385 train_values=y,
386 interpolated=bool(interpolate_step),
387 alpha=alpha,
388 warnings=notes,
389 )
392# --------------------------------------------------------------------------- #
393# endpoint prediction across a cohort
394# --------------------------------------------------------------------------- #
395def predict_endpoint(
396 scores: pd.DataFrame,
397 endpoints: dict[str, float] | None = None,
398 score_col: str = "relsa",
399 **kwargs,
400) -> list[Forecast]:
401 """Predict each animal's RELSA score at its humane endpoint.
403 This is the paper's primary evaluation: every measurement *up to the time
404 point immediately before* the endpoint trains the model, which then predicts
405 the score at the endpoint itself, where the actual score is known and can be
406 compared. ``endpoints`` maps animal id to endpoint time; omit it to use each
407 animal's last observed time point.
408 """
409 validate(scores, [score_col])
410 if endpoints is not None:
411 unknown = set(endpoints) - set(scores[ID_COL].astype(str))
412 if unknown or not np.isfinite(list(endpoints.values())).all():
413 raise ValueError("endpoint IDs must exist and endpoint times must be finite")
414 else:
415 warnings.warn("using last observed times as evaluation targets; these are not verified humane endpoints", stacklevel=2)
416 out: list[Forecast] = []
417 for animal, block in scores.groupby(ID_COL, sort=False):
418 block = block.dropna(subset=[score_col]).sort_values(TIME_COL)
419 if block.empty:
420 warnings.warn(f"animal {animal}: no finite RELSA scores", stacklevel=2)
421 continue
422 endpoint = (
423 float(block[TIME_COL].iloc[-1])
424 if endpoints is None
425 else endpoints.get(str(animal))
426 )
427 if endpoint is None:
428 continue
429 train = block[block[TIME_COL] < endpoint]
430 truth = block[block[TIME_COL] == endpoint]
431 if len(train) < 2:
432 warnings.warn(
433 f"animal {animal}: only {len(train)} points before the endpoint, skipping",
434 stacklevel=2,
435 )
436 continue
437 forecast = forecast_animal(
438 train[TIME_COL], train[score_col], target_times=[endpoint],
439 animal=str(animal), **kwargs,
440 )
441 if not truth.empty:
442 forecast.actual = np.array([float(truth[score_col].iloc[0])])
443 out.append(forecast)
444 return out
447def forecast_indirect(
448 prepared: pd.DataFrame,
449 reference,
450 target_time: float,
451 animal: str = "",
452 score_col: str = "relsa",
453 **kwargs,
454) -> dict[str, object]:
455 """The *indirect* prediction: forecast each variable, then score the forecasts.
457 ``prepared`` is the normalized measurement table for one animal (percent
458 scale) and ``reference`` a ``relsa_score.ReferenceModel``. Each variable is
459 forecast to ``target_time`` independently, and the RELSA score is computed
460 from those forecasts.
462 The paper found this worse than forecasting RELSA directly (median deviation
463 -0.240 vs -0.002, d = 1.42): each variable's forecast error propagates into
464 the score, whereas the direct forecast carries only its own error. Use this
465 to reproduce that comparison, not as the working method.
466 """
467 from relsa_score import relsa_scores
469 validate(prepared, list(reference.variables))
470 if prepared[ID_COL].nunique() != 1:
471 raise ValueError("indirect forecasting requires exactly one animal")
472 block = prepared.sort_values(TIME_COL)
473 animal = animal or str(block[ID_COL].iloc[0])
474 train = block[block[TIME_COL] < target_time]
475 row: dict[str, object] = {ID_COL: animal, TIME_COL: target_time}
476 predicted_row = {ID_COL: animal, TIME_COL: target_time}
477 for var in reference.variables:
478 series = train[[TIME_COL, var]].dropna()
479 if len(series) < 2:
480 predicted_row[var] = np.nan
481 continue
482 try:
483 forecast = forecast_animal(
484 series[TIME_COL], series[var], target_times=[target_time],
485 animal=f"{animal}:{var}", clip_at_zero=False, **kwargs,
486 )
487 predicted_row[var] = float(forecast.predicted[0])
488 except Exception as exc:
489 warnings.warn(f"indirect variable {var!r} failed: {exc}", stacklevel=2)
490 predicted_row[var] = np.nan
491 scored = relsa_scores(pd.DataFrame([predicted_row]), reference, keep_meta=False)
492 row["predicted"] = float(scored[score_col].iloc[0])
493 row.update({f"pred_{k}": v for k, v in predicted_row.items()
494 if k not in (ID_COL, TIME_COL)})
495 truth = block[np.isclose(block[TIME_COL], target_time)]
496 if not truth.empty:
497 actual = relsa_scores(truth, reference, keep_meta=False)
498 row["actual"] = float(actual[score_col].iloc[0])
499 return row
502def rolling_forecast(
503 times: Iterable[float],
504 values: Iterable[float],
505 min_train: int = 3,
506 animal: str = "",
507 **kwargs,
508) -> pd.DataFrame:
509 """One-step-ahead forecast at every time point, refitting as data accrue.
511 This is how the paper compares predictability across outcome measures
512 (Figure 2): at each time point, forecast the next one and record the
513 deviation from what actually happened.
514 """
515 t, y = _trajectory(times, values)
516 if min_train < 2:
517 raise ValueError("min_train must be at least two")
518 rows: list[dict[str, object]] = []
519 for cut in range(min_train, len(t)):
520 train_t, train_y = t[:cut], y[:cut]
521 if np.isfinite(train_y).sum() < 2:
522 continue
523 try:
524 forecast = forecast_animal(
525 train_t, train_y, target_times=[t[cut]], animal=animal, **kwargs
526 )
527 except Exception as exc: # a single failed refit must not stop the sweep
528 rows.append({ID_COL: animal, TIME_COL: t[cut], "predicted": np.nan,
529 "lower": np.nan, "upper": np.nan, "actual": y[cut],
530 "model": f"failed: {type(exc).__name__}"})
531 continue
532 rows.append({
533 ID_COL: animal,
534 TIME_COL: float(t[cut]),
535 "predicted": float(forecast.predicted[0]),
536 "lower": float(forecast.lower[0]),
537 "upper": float(forecast.upper[0]),
538 "actual": float(y[cut]),
539 "model": forecast.fit.label(),
540 "alpha": forecast.alpha,
541 "interpolated": forecast.interpolated,
542 "n_observed": len(forecast.train_times) if forecast.train_times is not None else cut,
543 "notes": "; ".join(forecast.warnings),
544 })
545 return pd.DataFrame(rows)
548def summarize(forecasts: Sequence[Forecast], group: dict[str, str] | None = None) -> pd.DataFrame:
549 """Per-animal and overall RMSE / PICP / MPIW, in the layout of the paper's Table 1."""
550 rows: list[dict[str, object]] = []
551 actual: list[float] = []
552 predicted: list[float] = []
553 lower: list[float] = []
554 upper: list[float] = []
555 for forecast in forecasts:
556 if forecast.actual is None:
557 continue
558 metrics = forecast.metrics()
559 rows.append({
560 "group": (group or {}).get(forecast.animal, ""),
561 ID_COL: forecast.animal,
562 "model": forecast.fit.label(),
563 "n": metrics.n, "n_interval": metrics.n_interval,
564 "rmse": round(metrics.rmse, 4),
565 "picp": round(metrics.picp, 1),
566 "mpiw": round(metrics.mpiw, 3),
567 })
568 actual.extend(forecast.actual.tolist())
569 predicted.extend(forecast.predicted.tolist())
570 lower.extend(forecast.lower.tolist())
571 upper.extend(forecast.upper.tolist())
573 frame = pd.DataFrame(rows)
574 if not actual:
575 return frame
576 overall = forecast_metrics(actual, predicted, lower, upper)
577 if not frame.empty and frame["group"].astype(bool).any():
578 for label, block in frame.groupby("group"):
579 if not label:
580 continue
581 sub = [f for f in forecasts if f.animal in set(block[ID_COL])]
582 agg = forecast_metrics(
583 np.concatenate([f.actual for f in sub]), # type: ignore[arg-type]
584 np.concatenate([f.predicted for f in sub]),
585 np.concatenate([f.lower for f in sub]),
586 np.concatenate([f.upper for f in sub]),
587 )
588 frame = pd.concat([frame, pd.DataFrame([{
589 "group": label, ID_COL: f"-- {label} --", "model": "",
590 "n": agg.n, "n_interval": agg.n_interval, "rmse": round(agg.rmse, 4),
591 "picp": round(agg.picp, 1), "mpiw": round(agg.mpiw, 3),
592 }])], ignore_index=True)
593 return pd.concat([frame, pd.DataFrame([{
594 "group": "", ID_COL: "OVERALL", "model": "",
595 "n": overall.n, "n_interval": overall.n_interval, "rmse": round(overall.rmse, 4),
596 "picp": round(overall.picp, 1), "mpiw": round(overall.mpiw, 3),
597 }])], ignore_index=True)
600def plot_forecast(
601 forecast: Forecast,
602 path: str | Path,
603 endpoint_threshold: float | None = None,
604 zones: Sequence[float] = (),
605 title: str | None = None,
606) -> Path:
607 """Trajectory, forecast, and interval in the style of the paper's Figure 1."""
608 import matplotlib
610 matplotlib.use("Agg")
611 import matplotlib.pyplot as plt
613 fig, ax = plt.subplots(figsize=(5.6, 3.8), constrained_layout=True)
614 if forecast.train_times is not None:
615 ax.plot(forecast.train_times, forecast.train_values, "-o", color="black",
616 ms=4, lw=1.3, label="actual RELSA")
617 ax.fill_between(
618 np.concatenate([[forecast.train_times[-1]], forecast.times]),
619 np.concatenate([[forecast.train_values[-1]], forecast.lower]),
620 np.concatenate([[forecast.train_values[-1]], forecast.upper]),
621 color="#4fc3c7", alpha=0.45,
622 label=f"{100 * (1 - forecast.alpha):.0f}% prediction interval",
623 )
624 ax.plot(forecast.times, forecast.predicted, "o--", color="#1f7a7d", ms=5,
625 label="predicted RELSA")
626 if forecast.actual is not None:
627 ax.plot(forecast.times, forecast.actual, "o", color="black", ms=5)
628 if endpoint_threshold is not None:
629 ax.axhline(endpoint_threshold, ls="--", color="#c1453b", lw=1.2)
630 ax.annotate("user-supplied comparison line", xy=(0.02, endpoint_threshold),
631 xycoords=("axes fraction", "data"), xytext=(0, 4),
632 textcoords="offset points", color="#c1453b", fontsize=8)
633 for zone in zones:
634 ax.axhline(zone, ls=":", color="#666666", lw=1.0)
635 ax.set_xlabel("time")
636 ax.set_ylabel("RELSA score")
637 ax.set_title(title or f"{forecast.animal} — {forecast.fit.label()}", fontsize=10)
638 ax.legend(frameon=False, fontsize=8, loc="lower right")
639 for spine in ("top", "right"):
640 ax.spines[spine].set_visible(False)
641 path = Path(path)
642 fig.savefig(path, dpi=200)
643 plt.close(fig)
644 return path
647# --------------------------------------------------------------------------- #
648# CLI
649# --------------------------------------------------------------------------- #
650def build_parser() -> argparse.ArgumentParser:
651 parser = argparse.ArgumentParser(
652 description="Explore RELSA forecasts with nonseasonal ARIMA (foRcast-style).",
653 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
654 )
655 parser.add_argument("scores", help="CSV of RELSA scores (output of relsa_score.py)")
656 parser.add_argument("--score-col", default="relsa")
657 parser.add_argument(
658 "--mode", choices=("endpoint", "rolling"), default="endpoint",
659 help="endpoint: predict the score at each animal's endpoint; "
660 "rolling: one-step-ahead forecast at every time point",
661 )
662 parser.add_argument(
663 "--endpoints", metavar="ID=TIME", action="append", default=[],
664 help="humane-endpoint time per animal, repeatable "
665 "(default: last-observation evaluation only, not a verified endpoint)",
666 )
667 parser.add_argument("--animals", help="comma-separated ids to restrict to")
668 parser.add_argument(
669 "--interpolate-step", type=float, default=DEFAULT_STEP,
670 help="grid for linear interpolation before fitting; 0 disables it",
671 )
672 parser.add_argument("--alpha", type=float, default=0.05,
673 help="1 - alpha is the prediction interval level")
674 parser.add_argument("--max-p", type=int, default=5)
675 parser.add_argument("--max-q", type=int, default=5)
676 parser.add_argument("--max-d", type=int, default=2)
677 parser.add_argument("--full-grid", action="store_true",
678 help="exhaustive p/q search instead of the stepwise hill-climb")
679 parser.add_argument("--group-col", help="column labelling model/intervention, for the summary")
680 parser.add_argument("--out", help="write per-forecast rows to this CSV")
681 parser.add_argument("--summary-out", help="write the metric summary to this CSV")
682 parser.add_argument("--plot-dir", help="write one figure per animal into this directory")
683 parser.add_argument("--endpoint-line", type=float, default=None,
684 help="user-supplied comparison line; not an automatically inferred endpoint")
685 return parser
688def main(argv: Sequence[str] | None = None) -> int:
689 args = build_parser().parse_args(argv)
690 frame = read_relsa_table(args.scores)
691 for col in (ID_COL, TIME_COL):
692 if col not in frame.columns:
693 raise SystemExit(f"column {col!r} missing from {args.scores}")
694 if args.score_col not in frame.columns:
695 raise SystemExit(f"score column {args.score_col!r} missing from {args.scores}")
696 validate(frame, [args.score_col])
698 keep = parse_list(args.animals)
699 if keep:
700 frame = frame[frame[ID_COL].astype(str).isin(keep)]
701 if frame.empty:
702 raise SystemExit(f"no rows for animals {keep}")
704 endpoints: dict[str, float] | None = None
705 if args.endpoints:
706 endpoints = {}
707 for pair in args.endpoints:
708 if "=" not in pair:
709 raise SystemExit(f"--endpoints expects ID=TIME, got {pair!r}")
710 animal, value = pair.split("=", 1)
711 endpoints[animal] = float(value)
713 groups = None
714 if args.group_col:
715 if args.group_col not in frame.columns:
716 raise SystemExit(f"--group-col {args.group_col!r} not in data")
717 groups = (
718 frame.groupby(ID_COL)[args.group_col].first().astype(str).to_dict()
719 )
720 groups = {str(k): v for k, v in groups.items()}
722 step = args.interpolate_step or None
723 arima_kwargs = dict(
724 interpolate_step=step,
725 alpha=args.alpha,
726 max_p=args.max_p,
727 max_q=args.max_q,
728 max_d=args.max_d,
729 stepwise=not args.full_grid,
730 )
732 if args.mode == "rolling":
733 rows = []
734 for animal, block in frame.groupby(ID_COL, sort=False):
735 block = block.dropna(subset=[args.score_col]).sort_values(TIME_COL)
736 rows.append(
737 rolling_forecast(
738 block[TIME_COL], block[args.score_col], animal=str(animal),
739 **arima_kwargs,
740 )
741 )
742 table = pd.concat(rows, ignore_index=True) if rows else pd.DataFrame()
743 if table.empty:
744 raise SystemExit("no forecasts could be produced")
745 metrics = forecast_metrics(
746 table["actual"], table["predicted"], table["lower"], table["upper"]
747 )
748 print(table.to_string(index=False))
749 print(
750 f"\nrolling one-step-ahead: n={metrics.n} n_interval={metrics.n_interval} RMSE={metrics.rmse:.4f} "
751 f"PICP={metrics.picp:.1f}% MPIW={metrics.mpiw:.3f}"
752 )
753 if args.out:
754 table.to_csv(args.out, index=False)
755 print(f"wrote {args.out}", file=sys.stderr)
756 return 0
758 forecasts = predict_endpoint(
759 frame, endpoints=endpoints, score_col=args.score_col, **arima_kwargs
760 )
761 if not forecasts:
762 raise SystemExit("no animal had enough data to forecast")
764 detail = pd.concat([f.to_frame() for f in forecasts], ignore_index=True)
765 print(detail.to_string(index=False))
766 summary = summarize(forecasts, group=groups)
767 print("\n" + summary.to_string(index=False))
769 for forecast in forecasts:
770 for note in forecast.warnings:
771 print(f"note [{forecast.animal}]: {note}", file=sys.stderr)
773 if args.out:
774 detail.to_csv(args.out, index=False)
775 print(f"wrote {args.out}", file=sys.stderr)
776 if args.summary_out:
777 summary.to_csv(args.summary_out, index=False)
778 print(f"wrote {args.summary_out}", file=sys.stderr)
779 if args.plot_dir:
780 directory = Path(args.plot_dir)
781 directory.mkdir(parents=True, exist_ok=True)
782 for index, forecast in enumerate(forecasts, 1):
783 safe_id = "".join(c if c.isalnum() or c in "-_." else "_" for c in forecast.animal)
784 written = plot_forecast(
785 forecast, directory / f"{index}-{safe_id}.png",
786 endpoint_threshold=args.endpoint_line,
787 )
788 print(f"wrote {written}", file=sys.stderr)
789 return 0
792if __name__ == "__main__":
793 raise SystemExit(main())