christiangeorgelucas/time-series-tools

from statsmodels.tsa.stattools import acf as sm_acf

from gen.messages_pb2 import AcfInput, CorrelogramResult
from gen.axiom_context import AxiomContext
from nodes._common import err, to_pylist, validate_series


def acf(ax: AxiomContext, input: AcfInput) -> CorrelogramResult:
    """Compute the sample autocorrelation function of a series at lags
    0..nlags, optionally with Bartlett confidence intervals, wrapping
    statsmodels.tsa.stattools.acf.
    """
    values = list(input.values)
    e = validate_series(values, min_len=2)
    if e:
        return CorrelogramResult(error=e)

    nlags = input.nlags if input.nlags > 0 else None
    if nlags is not None and nlags >= len(values):
        return CorrelogramResult(error=err("INVALID_ARGUMENT", f"nlags ({nlags}) must be less than the series length ({len(values)})"))

    alpha = input.alpha if input.alpha > 0 else None
    if alpha is not None and not (0 < alpha < 1):
        return CorrelogramResult(error=err("INVALID_ARGUMENT", f"alpha must be in (0, 1), got {alpha}"))

    try:
        if alpha is not None:
            vals, confint = sm_acf(values, nlags=nlags, fft=input.fft, alpha=alpha)
        else:
            vals = sm_acf(values, nlags=nlags, fft=input.fft, alpha=None)
            confint = None
    except Exception as exc:
        return CorrelogramResult(error=err("COMPUTE_ERROR", str(exc)))

    result = CorrelogramResult(values=to_pylist(vals))
    if confint is not None:
        result.conf_lower.extend(float(c[0]) for c in confint)
        result.conf_upper.extend(float(c[1]) for c in confint)
    return result

from statsmodels.tsa.stattools import adfuller

from gen.messages_pb2 import AdfInput, AdfResult
from gen.axiom_context import AxiomContext
from nodes._common import err, validate_series


def adf_test(ax: AxiomContext, input: AdfInput) -> AdfResult:
    """Run the Augmented Dickey-Fuller test for a unit root — the null
    hypothesis is that the series is NON-stationary — wrapping
    statsmodels.tsa.stattools.adfuller.
    """
    values = list(input.values)
    e = validate_series(values, min_len=4)
    if e:
        return AdfResult(error=e)

    regression = (input.regression or "c").lower()
    if regression not in ("c", "ct", "ctt", "n"):
        return AdfResult(error=err("INVALID_ARGUMENT", f"unknown regression '{input.regression}', expected 'c', 'ct', 'ctt', or 'n'"))

    autolag_raw = input.autolag or "AIC"
    autolag = None if autolag_raw.lower() == "none" else autolag_raw
    if autolag is not None and autolag.upper() not in ("AIC", "BIC", "T-STAT"):
        return AdfResult(error=err("INVALID_ARGUMENT", f"unknown autolag '{input.autolag}', expected 'AIC', 'BIC', 't-stat', or 'none'"))

    maxlag = input.maxlag if input.HasField("maxlag") else None
    if maxlag is not None and maxlag < 0:
        return AdfResult(error=err("INVALID_ARGUMENT", f"maxlag must be >= 0, got {maxlag}"))
    if maxlag is not None and maxlag >= len(values) // 2:
        return AdfResult(error=err("INVALID_ARGUMENT", f"maxlag ({maxlag}) too large for a series of length {len(values)}"))

    try:
        stat, pvalue, used_lag, nobs, crit, *_ = adfuller(values, maxlag=maxlag, regression=regression, autolag=autolag)
    except Exception as exc:
        return AdfResult(error=err("COMPUTE_ERROR", str(exc)))

    return AdfResult(
        statistic=float(stat),
        p_value=float(pvalue),
        used_lag=int(used_lag),
        nobs=int(nobs),
        critical_value_1pct=float(crit.get("1%", float("nan"))),
        critical_value_5pct=float(crit.get("5%", float("nan"))),
        critical_value_10pct=float(crit.get("10%", float("nan"))),
        stationary=bool(pvalue < 0.05),
    )

import numpy as np
from statsmodels.tsa.arima.model import ARIMA

from gen.messages_pb2 import ArimaInput, ForecastResult
from gen.axiom_context import AxiomContext
from nodes._common import MAX_DIFF_ORDER, MAX_ORDER, err, to_pylist, validate_series

_VALID_TRENDS = {"n", "c", "t", "ct"}


def arima_forecast(ax: AxiomContext, input: ArimaInput) -> ForecastResult:
    """Fit a (Seasonal) ARIMA model by maximum likelihood and forecast
    forward `horizon` steps, wrapping statsmodels.tsa.arima.model.ARIMA.
    """
    values = list(input.values)
    e = validate_series(values, min_len=8)
    if e:
        return ForecastResult(error=e)

    p, d, q = input.p, input.d, input.q
    sp, sd, sq, sper = input.seasonal_p, input.seasonal_d, input.seasonal_q, input.seasonal_period

    for name, v, cap in (
        ("p", p, MAX_ORDER), ("q", q, MAX_ORDER), ("d", d, MAX_DIFF_ORDER),
        ("seasonal_p", sp, MAX_ORDER), ("seasonal_q", sq, MAX_ORDER), ("seasonal_d", sd, MAX_DIFF_ORDER),
    ):
        if v < 0 or v > cap:
            return ForecastResult(error=err("INVALID_ARGUMENT", f"{name} must be in [0, {cap}], got {v}"))
    if sper < 0:
        return ForecastResult(error=err("INVALID_ARGUMENT", f"seasonal_period out of range: {sper}"))
    if (sp or sd or sq) and sper < 2:
        return ForecastResult(error=err("INVALID_ARGUMENT", "seasonal_period must be >= 2 when a seasonal order is set"))

    horizon = input.horizon if input.horizon > 0 else 1

    trend = input.trend or None
    if trend is not None and trend not in _VALID_TRENDS:
        return ForecastResult(error=err("INVALID_ARGUMENT", f"unknown trend '{input.trend}', expected one of {sorted(_VALID_TRENDS)}"))

    seasonal_order = (sp, sd, sq, sper) if sper >= 2 else (0, 0, 0, 0)

    try:
        model = ARIMA(values, order=(p, d, q), seasonal_order=seasonal_order, trend=trend)
        res = model.fit()
        fc = res.get_forecast(steps=horizon)
        mean = np.asarray(fc.predicted_mean, dtype=float)
        ci = np.asarray(fc.conf_int(alpha=0.05), dtype=float)
        fitted = np.asarray(res.fittedvalues, dtype=float)
        resid = np.asarray(res.resid, dtype=float)
    except Exception as exc:
        return ForecastResult(error=err("COMPUTE_ERROR", str(exc)))

    return ForecastResult(
        forecast=to_pylist(mean),
        fitted=to_pylist(fitted),
        aic=float(res.aic),
        bic=float(res.bic),
        sse=float(np.nansum(resid ** 2)),
        conf_lower=to_pylist(ci[:, 0]),
        conf_upper=to_pylist(ci[:, 1]),
    )

import numpy as np
from statsmodels.tsa.seasonal import STL, seasonal_decompose

from gen.messages_pb2 import DecomposeInput, DecomposeResult
from gen.axiom_context import AxiomContext
from nodes._common import err, to_pylist, validate_series


def decompose(ax: AxiomContext, input: DecomposeInput) -> DecomposeResult:
    """Split a series into trend, seasonal, and residual components via STL
    (Seasonal-Trend decomposition using LOESS) or classical moving-average
    decomposition, wrapping statsmodels.tsa.seasonal.STL / seasonal_decompose.
    Needs at least two full seasonal periods of data.
    """
    values = list(input.values)
    e = validate_series(values, min_len=4)
    if e:
        return DecomposeResult(error=e)

    method = (input.method or "stl").lower()
    if method not in ("stl", "classical"):
        return DecomposeResult(error=err("INVALID_ARGUMENT", f"unknown method '{input.method}', expected 'stl' or 'classical'"))

    period = input.period
    if period < 2:
        return DecomposeResult(error=err("INVALID_ARGUMENT", f"period must be >= 2, got {period}"))
    if len(values) < 2 * period:
        return DecomposeResult(error=err("INVALID_INPUT", f"need at least 2*period ({2 * period}) observations, got {len(values)}"))

    arr = np.asarray(values, dtype=float)

    try:
        if method == "stl":
            res = STL(arr, period=period, robust=input.robust).fit()
            observed, trend, seasonal, resid = arr, res.trend, res.seasonal, res.resid
        else:
            model = (input.model or "additive").lower()
            if model not in ("additive", "multiplicative"):
                return DecomposeResult(error=err("INVALID_ARGUMENT", f"unknown model '{input.model}', expected 'additive' or 'multiplicative'"))
            if model == "multiplicative" and np.any(arr <= 0):
                return DecomposeResult(error=err("INVALID_ARGUMENT", "multiplicative decomposition requires all values > 0"))
            res = seasonal_decompose(arr, model=model, period=period)
            observed, trend, seasonal, resid = res.observed, res.trend, res.seasonal, res.resid
    except Exception as exc:
        return DecomposeResult(error=err("COMPUTE_ERROR", str(exc)))

    return DecomposeResult(
        observed=to_pylist(observed),
        trend=to_pylist(trend),
        seasonal=to_pylist(seasonal),
        resid=to_pylist(resid),
    )

import numpy as np
from statsmodels.tsa.tsatools import detrend as sm_detrend

from gen.messages_pb2 import DetrendInput, Series
from gen.axiom_context import AxiomContext
from nodes._common import err, to_pylist, validate_series


def detrend(ax: AxiomContext, input: DetrendInput) -> Series:
    """Remove a constant (the sample mean) or best-fit linear trend from a
    series, wrapping statsmodels.tsa.tsatools.detrend. Emits the canonical
    Series envelope, so it chains directly into a second Difference/Detrend
    call or into any other node's `values` field.
    """
    values = list(input.values)
    e = validate_series(values, min_len=2)
    if e:
        return Series(error=e)

    method = (input.method or "linear").lower()
    if method not in ("linear", "constant"):
        return Series(error=err("INVALID_ARGUMENT", f"unknown method '{input.method}', expected 'linear' or 'constant'"))
    order = 1 if method == "linear" else 0

    arr = np.asarray(values, dtype=float)
    try:
        out = sm_detrend(arr, order=order)
    except Exception as exc:
        return Series(error=err("COMPUTE_ERROR", str(exc)))

    return Series(values=to_pylist(out))

import numpy as np

from gen.messages_pb2 import DifferenceInput, Series
from gen.axiom_context import AxiomContext
from nodes._common import MAX_DIFF_ORDER, err, to_pylist, validate_series


def _diff_once(arr: np.ndarray, lag: int) -> np.ndarray:
    return arr[lag:] - arr[:-lag]


def difference(ax: AxiomContext, input: DifferenceInput) -> Series:
    """Apply non-seasonal and/or seasonal differencing to a series — the
    standard transform for turning a non-stationary series into a stationary
    one before modeling. Seasonal differencing (lag = seasonal_periods) is
    applied first, then non-seasonal (lag 1) differencing; since both are
    linear filters the two commute, so this ordering does not affect the
    result. Emits the canonical Series envelope, so it chains directly into
    a second Difference/Detrend call or into any other node's `values`
    field.
    """
    values = list(input.values)
    e = validate_series(values, min_len=2)
    if e:
        return Series(error=e)

    order = input.order if input.HasField("order") else 1
    if order < 0:
        return Series(error=err("INVALID_ARGUMENT", f"order must be >= 0, got {order}"))
    if order > MAX_DIFF_ORDER:
        return Series(error=err("LIMIT_EXCEEDED", f"order {order} exceeds the {MAX_DIFF_ORDER} limit"))

    seasonal_periods = input.seasonal_periods
    seasonal_order = 0
    if seasonal_periods > 0:
        seasonal_order = input.seasonal_order if input.seasonal_order > 0 else 1
        if seasonal_order > MAX_DIFF_ORDER:
            return Series(error=err("LIMIT_EXCEEDED", f"seasonal_order {seasonal_order} exceeds the {MAX_DIFF_ORDER} limit"))
    elif input.seasonal_order > 0:
        return Series(error=err("INVALID_ARGUMENT", "seasonal_order was set but seasonal_periods was not"))

    min_needed = order + seasonal_order * max(seasonal_periods, 1) + 1
    if len(values) < min_needed:
        return Series(error=err("INVALID_INPUT", f"need at least {min_needed} observations for this differencing configuration, got {len(values)}"))

    out = np.asarray(values, dtype=float)
    try:
        for _ in range(seasonal_order):
            out = _diff_once(out, seasonal_periods)
        for _ in range(order):
            out = _diff_once(out, 1)
    except Exception as exc:
        return Series(error=err("COMPUTE_ERROR", str(exc)))

    return Series(values=to_pylist(out))

import numpy as np
from statsmodels.tsa.holtwinters import ExponentialSmoothing

from gen.messages_pb2 import EtsInput, ForecastResult
from gen.axiom_context import AxiomContext
from nodes._common import err, to_pylist, validate_series

_TREND_MAP = {"none": None, "add": "add", "mul": "mul"}
_SEASONAL_MAP = {"none": None, "add": "add", "mul": "mul"}


def exponential_smoothing_forecast(ax: AxiomContext, input: EtsInput) -> ForecastResult:
    """Fit a Holt-Winters exponential smoothing model — optional additive or
    multiplicative trend (with damping) and seasonal components — by
    minimizing sum-of-squared-errors, and forecast forward `horizon` steps,
    wrapping statsmodels.tsa.holtwinters.ExponentialSmoothing.
    """
    values = list(input.values)
    e = validate_series(values, min_len=4)
    if e:
        return ForecastResult(error=e)

    trend_key = (input.trend or "none").lower()
    seasonal_key = (input.seasonal or "none").lower()
    if trend_key not in _TREND_MAP:
        return ForecastResult(error=err("INVALID_ARGUMENT", f"unknown trend '{input.trend}', expected 'none', 'add', or 'mul'"))
    if seasonal_key not in _SEASONAL_MAP:
        return ForecastResult(error=err("INVALID_ARGUMENT", f"unknown seasonal '{input.seasonal}', expected 'none', 'add', or 'mul'"))

    trend = _TREND_MAP[trend_key]
    seasonal = _SEASONAL_MAP[seasonal_key]
    seasonal_periods = input.seasonal_periods if seasonal else None

    if seasonal and seasonal_periods < 2:
        return ForecastResult(error=err("INVALID_ARGUMENT", "seasonal_periods must be >= 2 when seasonal is set"))
    if seasonal and len(values) < 2 * seasonal_periods:
        return ForecastResult(error=err("INVALID_INPUT", f"need at least 2*seasonal_periods ({2 * seasonal_periods}) observations, got {len(values)}"))

    arr = np.asarray(values, dtype=float)
    if seasonal == "mul" and np.any(arr <= 0):
        return ForecastResult(error=err("INVALID_ARGUMENT", "multiplicative seasonal component requires all values > 0"))
    if trend == "mul" and np.any(arr <= 0):
        return ForecastResult(error=err("INVALID_ARGUMENT", "multiplicative trend requires all values > 0"))

    horizon = input.horizon if input.horizon > 0 else 1

    damped = bool(input.damped_trend) and trend is not None

    try:
        model = ExponentialSmoothing(
            arr, trend=trend, seasonal=seasonal, seasonal_periods=seasonal_periods,
            damped_trend=damped, initialization_method="estimated",
        )
        res = model.fit()
        forecast = np.asarray(res.forecast(horizon), dtype=float)
        fitted = np.asarray(res.fittedvalues, dtype=float)
    except Exception as exc:
        return ForecastResult(error=err("COMPUTE_ERROR", str(exc)))

    return ForecastResult(
        forecast=to_pylist(forecast),
        fitted=to_pylist(fitted),
        aic=float(res.aic),
        bic=float(res.bic),
        sse=float(res.sse),
    )

import contextlib
import io

import numpy as np
from statsmodels.tsa.stattools import grangercausalitytests

from gen.messages_pb2 import GrangerInput, GrangerResult
from gen.axiom_context import AxiomContext
from nodes._common import MAX_ORDER, err, validate_series

_MAX_MAXLAG = MAX_ORDER * 2


def granger_causality(ax: AxiomContext, input: GrangerInput) -> GrangerResult:
    """Test whether the past of `predictor` helps forecast `target` beyond
    target's own past, at every lag from 1 to maxlag, wrapping
    statsmodels.tsa.stattools.grangercausalitytests. `target` and `predictor`
    must be the same length and time-aligned. A statistical precedence test,
    not proof of causation.
    """
    target = list(input.target)
    predictor = list(input.predictor)

    e = validate_series(target, min_len=4, name="target")
    if e:
        return GrangerResult(error=e)
    e = validate_series(predictor, min_len=4, name="predictor")
    if e:
        return GrangerResult(error=e)
    if len(target) != len(predictor):
        return GrangerResult(error=err("INVALID_INPUT", f"target ({len(target)}) and predictor ({len(predictor)}) must be the same length"))

    maxlag = input.maxlag
    if maxlag < 1:
        return GrangerResult(error=err("INVALID_ARGUMENT", f"maxlag must be >= 1, got {maxlag}"))
    if maxlag > _MAX_MAXLAG:
        return GrangerResult(error=err("LIMIT_EXCEEDED", f"maxlag {maxlag} exceeds the {_MAX_MAXLAG} limit"))

    n = len(target)
    if maxlag >= n // 2:
        return GrangerResult(error=err("INVALID_ARGUMENT", f"maxlag ({maxlag}) too large for a series of length {n}"))

    addconst = input.addconst if input.HasField("addconst") else True

    arr = np.column_stack([target, predictor])
    buf = io.StringIO()
    try:
        with contextlib.redirect_stdout(buf):
            results = grangercausalitytests(arr, maxlag=maxlag, addconst=addconst)
    except Exception as exc:
        return GrangerResult(error=err("COMPUTE_ERROR", str(exc)))

    lags_out, f_stat, f_pval, chi2_pval, lr_pval = [], [], [], [], []
    best_lag, best_p = 0, None
    for lag in range(1, maxlag + 1):
        test = results[lag][0]
        fstat, fpval, _, _ = test["ssr_ftest"]
        _, chi2p, _ = test["ssr_chi2test"]
        _, lrp, _ = test["lrtest"]
        lags_out.append(lag)
        f_stat.append(float(fstat))
        f_pval.append(float(fpval))
        chi2_pval.append(float(chi2p))
        lr_pval.append(float(lrp))
        if best_p is None or fpval < best_p:
            best_p, best_lag = fpval, lag

    return GrangerResult(
        lags=lags_out,
        ssr_ftest_stat=f_stat,
        ssr_ftest_pvalue=f_pval,
        ssr_chi2_pvalue=chi2_pval,
        lr_pvalue=lr_pval,
        best_lag=best_lag,
    )

import warnings

from statsmodels.tsa.stattools import kpss

from gen.messages_pb2 import KpssInput, KpssResult
from gen.axiom_context import AxiomContext
from nodes._common import err, validate_series


def kpss_test(ax: AxiomContext, input: KpssInput) -> KpssResult:
    """Run the KPSS test for stationarity — the null hypothesis is that the
    series IS stationary, the opposite null from AdfTest — wrapping
    statsmodels.tsa.stattools.kpss.
    """
    values = list(input.values)
    e = validate_series(values, min_len=4)
    if e:
        return KpssResult(error=e)

    regression = (input.regression or "c").lower()
    if regression not in ("c", "ct"):
        return KpssResult(error=err("INVALID_ARGUMENT", f"unknown regression '{input.regression}', expected 'c' or 'ct'"))

    nlags = "auto" if input.lags <= 0 else input.lags
    if isinstance(nlags, int) and nlags >= len(values):
        return KpssResult(error=err("INVALID_ARGUMENT", f"lags ({nlags}) must be less than the series length ({len(values)})"))

    try:
        # statsmodels warns when the statistic falls outside its tabulated
        # p-value range and clips the p-value instead — that clipped p-value
        # (still meaningful, just an interval bound) is exactly what we
        # return, so the warning is expected noise, not a signal to surface.
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            stat, pvalue, lags_used, crit = kpss(values, regression=regression, nlags=nlags)
    except Exception as exc:
        return KpssResult(error=err("COMPUTE_ERROR", str(exc)))

    return KpssResult(
        statistic=float(stat),
        p_value=float(pvalue),
        lags_used=int(lags_used),
        critical_value_10pct=float(crit.get("10%", float("nan"))),
        critical_value_5pct=float(crit.get("5%", float("nan"))),
        critical_value_2_5pct=float(crit.get("2.5%", float("nan"))),
        critical_value_1pct=float(crit.get("1%", float("nan"))),
        stationary=bool(pvalue >= 0.05),
    )

from statsmodels.stats.diagnostic import acorr_ljungbox

from gen.messages_pb2 import LjungBoxInput, LjungBoxResult
from gen.axiom_context import AxiomContext
from nodes._common import err, to_pylist, validate_series


def ljung_box_test(ax: AxiomContext, input: LjungBoxInput) -> LjungBoxResult:
    """Test whether a group of autocorrelations (typically a model's
    residuals) are jointly zero — i.e. whether the series looks like white
    noise — at one or more lags, optionally also reporting the older
    Box-Pierce statistic, wrapping
    statsmodels.stats.diagnostic.acorr_ljungbox.
    """
    values = list(input.values)
    e = validate_series(values, min_len=2)
    if e:
        return LjungBoxResult(error=e)

    lags = list(input.lags) if len(input.lags) > 0 else None
    if lags is not None:
        for lag in lags:
            if lag < 1 or lag >= len(values):
                return LjungBoxResult(error=err("INVALID_ARGUMENT", f"lag {lag} out of range [1, {len(values) - 1}]"))
    if input.model_df < 0:
        return LjungBoxResult(error=err("INVALID_ARGUMENT", f"model_df must be >= 0, got {input.model_df}"))

    try:
        df = acorr_ljungbox(values, lags=lags, boxpierce=input.box_pierce, model_df=input.model_df)
    except Exception as exc:
        return LjungBoxResult(error=err("COMPUTE_ERROR", str(exc)))

    result = LjungBoxResult(
        lags=[int(x) for x in df.index.tolist()],
        lb_stat=to_pylist(df["lb_stat"].to_numpy()),
        lb_pvalue=to_pylist(df["lb_pvalue"].to_numpy()),
    )
    if input.box_pierce:
        result.bp_stat.extend(to_pylist(df["bp_stat"].to_numpy()))
        result.bp_pvalue.extend(to_pylist(df["bp_pvalue"].to_numpy()))
    return result

from statsmodels.tsa.stattools import pacf as sm_pacf

from gen.messages_pb2 import PacfInput, CorrelogramResult
from gen.axiom_context import AxiomContext
from nodes._common import err, to_pylist, validate_series

_VALID_METHODS = {"ywadjusted", "ywm", "ols", "ld"}


def pacf(ax: AxiomContext, input: PacfInput) -> CorrelogramResult:
    """Compute the sample partial autocorrelation function of a series at
    lags 0..nlags, optionally with confidence intervals, wrapping
    statsmodels.tsa.stattools.pacf.
    """
    values = list(input.values)
    e = validate_series(values, min_len=3)
    if e:
        return CorrelogramResult(error=e)

    method = (input.method or "ywadjusted").lower()
    if method not in _VALID_METHODS:
        return CorrelogramResult(error=err("INVALID_ARGUMENT", f"unknown method '{input.method}', expected one of {sorted(_VALID_METHODS)}"))

    n = len(values)
    nlags = input.nlags if input.nlags > 0 else None
    if nlags is not None and nlags >= n:
        return CorrelogramResult(error=err("INVALID_ARGUMENT", f"nlags ({nlags}) must be less than the series length ({n})"))

    alpha = input.alpha if input.alpha > 0 else None
    if alpha is not None and not (0 < alpha < 1):
        return CorrelogramResult(error=err("INVALID_ARGUMENT", f"alpha must be in (0, 1), got {alpha}"))

    try:
        if alpha is not None:
            vals, confint = sm_pacf(values, nlags=nlags, method=method, alpha=alpha)
        else:
            vals = sm_pacf(values, nlags=nlags, method=method, alpha=None)
            confint = None
    except Exception as exc:
        return CorrelogramResult(error=err("COMPUTE_ERROR", str(exc)))

    result = CorrelogramResult(values=to_pylist(vals))
    if confint is not None:
        result.conf_lower.extend(float(c[0]) for c in confint)
        result.conf_upper.extend(float(c[1]) for c in confint)
    return result