diff --git a/README.md b/README.md index ea091327..5d5f802a 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ NNS is built around partial moments, the lower and upper components of variance, |---|---| | Distribution package | `ovvo-nns` | | Import package | `nns` | -| Current version | `1.0.9` | +| Current version | `1.1.0` | | Python | `>=3.11` | | Required runtime dependencies | NumPy, SciPy | | R required at runtime | No | diff --git a/pyproject.toml b/pyproject.toml index 8bea9a28..0ca517c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "ovvo-nns" -version = "1.0.9" +version = "1.1.0" description = "Python port of nonlinear nonparametric statistics from R NNS" readme = "README.md" requires-python = ">=3.11" diff --git a/src/nns/__init__.py b/src/nns/__init__.py index c4654ec5..848fc553 100644 --- a/src/nns/__init__.py +++ b/src/nns/__init__.py @@ -4,10 +4,20 @@ from nns.pm_matrix import pm_matrix as pm_matrix -__version__ = "1.0.9" +__version__ = "1.1.0" _EXPORTS = { + "BoostResult": ("nns.boost", "BoostResult"), "FactorDesign": ("nns.regression", "FactorDesign"), + "MRegResult": ("nns.multivariate_regression", "MRegResult"), + "MebootResult": ("nns.meboot", "MebootResult"), + "PartResult": ("nns.part", "PartResult"), + "RegFitted": ("nns.regression", "RegFitted"), + "RegPoints": ("nns.regression", "RegPoints"), + "RegResult": ("nns.regression", "RegResult"), + "SeasonalityResult": ("nns.seasonality", "SeasonalityResult"), + "StackResult": ("nns.stack", "StackResult"), + "VarResult": ("nns.var", "VarResult"), "causal_matrix": ("nns.causation", "causal_matrix"), "co_lpm": ("nns.co_moments", "co_lpm"), "co_lpm_nd": ("nns.dependence", "co_lpm_nd"), diff --git a/src/nns/boost.py b/src/nns/boost.py index a35f86d5..99e44085 100644 --- a/src/nns/boost.py +++ b/src/nns/boost.py @@ -4,7 +4,7 @@ import math import warnings from collections.abc import Callable, Sequence -from typing import Any, Literal, cast +from typing import Any, Literal, NotRequired, TypedDict, cast import numpy as np from numpy.typing import NDArray @@ -13,6 +13,8 @@ from nns.dependence import _gravity from nns.regression import ( Order, + RegResult, + RegXStar, _normalize_type, _prepare_y_values, _r_minmax_columns, @@ -22,7 +24,16 @@ from nns.stack import nns_stack Objective = Literal["min", "max"] -BoostResult = dict[str, Any] +BoostResult = TypedDict( + "BoostResult", + { + "feature.weights": NDArray[np.float64], + "feature.frequency": NDArray[np.float64], + "results": NotRequired[NDArray[np.float64]], + "pred.int": NotRequired["dict[str, NDArray[np.float64]] | None"], + }, +) +"""``nns_boost`` result. ``features_only=True`` returns only the feature keys.""" def nns_boost( @@ -267,7 +278,8 @@ def _nns_boost_core( order=depth, point_only=False, ) - xstar_train = np.asarray(xstar_fit["x.star"]["x"], dtype=np.float64) + x_star = cast("RegXStar", cast("RegResult", xstar_fit)["x.star"]) + xstar_train = np.asarray(x_star["x"], dtype=np.float64) xstar_train = _fill_nan_with_gravity(xstar_train) xstar_test = _project_xstar(x_train, x_test, coef) xstar_test = _fill_nan_with_gravity(xstar_test) diff --git a/src/nns/cdf.py b/src/nns/cdf.py index fd0d803d..c1ca49aa 100644 --- a/src/nns/cdf.py +++ b/src/nns/cdf.py @@ -2,13 +2,16 @@ import math from collections.abc import Sequence -from typing import Any, cast +from typing import TYPE_CHECKING, cast import numpy as np from numpy.typing import NDArray from nns.dependence import co_lpm_nd -from nns.regression import nns_reg +from nns.regression import RegResult, nns_reg + +if TYPE_CHECKING: + from nns.multivariate_regression import MRegResult def nns_cdf( @@ -73,7 +76,7 @@ def _univariate_cdf( }[type_value] y = pval.copy() - fit: dict[str, Any] | None = None + fit: RegResult | MRegResult | None = None if type_value == "survival": y = 1.0 - y elif type_value == "hazard": diff --git a/src/nns/meboot.py b/src/nns/meboot.py index d16a012b..4c194431 100644 --- a/src/nns/meboot.py +++ b/src/nns/meboot.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any +from typing import NotRequired, TypedDict import numpy as np from numpy.typing import NDArray @@ -8,7 +8,26 @@ from nns._helpers import _fast_lm from nns.dependence import nns_dep -MebootResult = dict[str, NDArray[np.float64] | float | None] + +class MebootResult(TypedDict): + """``nns_meboot`` per-``rho`` result following R's NNS.meboot diagnostics. + + Degenerate inputs return partial dictionaries (single-observation input + yields only ``x``), hence every key is optional. + """ + + x: NotRequired[NDArray[np.float64]] + replicates: NotRequired[NDArray[np.float64]] + ensemble: NotRequired[NDArray[np.float64]] + xx: NotRequired[NDArray[np.float64]] + z: NotRequired[NDArray[np.float64]] + dv: NotRequired[NDArray[np.float64]] + dvtrim: NotRequired[float] + xmin: NotRequired[float] + xmax: NotRequired[float] + desintxb: NotRequired[NDArray[np.float64]] + ordxx: NotRequired[NDArray[np.float64]] + kappa: NotRequired[float | None] def nns_meboot( @@ -30,7 +49,7 @@ def nns_meboot( elaps: bool = False, digits: int = 6, random_seed: int | None = None, -) -> dict[str, Any] | list[dict[str, Any]]: +) -> MebootResult | list[MebootResult]: """Maximum-entropy bootstrap matching R's NNS.meboot structure. Stochastic draws use NumPy's RNG, so exact replicate parity with R is not @@ -118,7 +137,7 @@ def _nns_meboot_one( sym: bool, digits: int, rng: np.random.Generator, -) -> dict[str, Any]: +) -> MebootResult: n = x.size time = np.arange(1, n + 1, dtype=np.float64) intercept, orig_drift = _fast_lm(time, x) diff --git a/src/nns/multivariate_regression.py b/src/nns/multivariate_regression.py index dcb66b74..74926961 100644 --- a/src/nns/multivariate_regression.py +++ b/src/nns/multivariate_regression.py @@ -1,7 +1,7 @@ from __future__ import annotations import math -from typing import Any, Literal, cast +from typing import Any, Literal, NotRequired, TypedDict, cast import numpy as np from numpy.typing import NDArray @@ -16,7 +16,38 @@ from nns.var import upm_var NBest = int | Literal["all"] | None -MRegResult = dict[str, Any] + +MRegFitted = dict[str, NDArray[np.float64] | NDArray[np.str_]] +"""Fit table keyed ``V1..Vn`` per regressor plus ``y``, ``y.hat``, ``NNS.ID``, +``residuals``, and (with ``confidence_interval``) ``conf.int.pos``/``conf.int.neg``. +Column count depends on the input, so keys stay dynamic.""" + +MRegPredInt = TypedDict( + "MRegPredInt", + { + "lower.pred.int": NDArray[np.float64], + "upper.pred.int": NDArray[np.float64], + }, +) +"""Prediction interval bounds for the requested ``point_est`` rows.""" + +MRegResult = TypedDict( + "MRegResult", + { + "Point.est": "NDArray[np.float64] | None", + "RPM": dict[str, NDArray[np.float64]], + "R2": NotRequired[float], + "rhs.partitions": NotRequired[dict[str, NDArray[np.float64]]], + "pred.int": NotRequired["MRegPredInt | None"], + "Fitted.xy": NotRequired[MRegFitted], + }, +) +"""``nns_m_reg`` result. + +``point_only=True`` returns only ``Point.est`` and ``RPM``; the full form always +carries the remaining keys (``pred.int`` is ``None`` unless +``confidence_interval`` is set). +""" def nns_m_reg( @@ -482,7 +513,7 @@ def _apply_multivariate_intervals( point_predictions: NDArray[np.float64] | None, *, confidence_interval: float | None, -) -> dict[str, NDArray[np.float64]] | None: +) -> MRegPredInt | None: if confidence_interval is None: return None alpha = (1.0 - float(confidence_interval)) / 2.0 diff --git a/src/nns/regression.py b/src/nns/regression.py index 9027f765..389b6f68 100644 --- a/src/nns/regression.py +++ b/src/nns/regression.py @@ -1,9 +1,9 @@ from __future__ import annotations import math -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import Any, Literal, cast +from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict, cast, overload import numpy as np from numpy.typing import NDArray @@ -15,12 +15,92 @@ from nns.central_tendencies import nns_mode from nns.copula import _copula from nns.dependence import _gravity, nns_dep -from nns.part import NoiseReduction, nns_part +from nns.part import NoiseReduction, PartResult, nns_part from nns.smoothing import r_smooth_spline_fixed_spar from nns.var import lpm_var, upm_var Order = int | Literal["max"] | None +if TYPE_CHECKING: + from nns.multivariate_regression import MRegResult + +# Result shapes mirror installed R NNS.reg output names, including dotted keys, +# hence the functional TypedDict syntax. + +class RegPoints(TypedDict): + x: NDArray[np.float64] + y: NDArray[np.float64] +"""Consolidated regression points; also the whole result when ``multivariate_call=True``.""" + +DerivativeTable = TypedDict( + "DerivativeTable", + { + "Coefficient": NDArray[np.float64], + "X.Lower.Range": NDArray[np.float64], + "X.Upper.Range": NDArray[np.float64], + }, +) +"""Piecewise slope per x-interval, as in R's ``$derivative`` table.""" + +RegPredInt = TypedDict( + "RegPredInt", + { + "pred.int.neg": NDArray[np.float64], + "pred.int.pos": NDArray[np.float64], + }, +) +"""Prediction interval bounds for the requested ``point_est`` rows.""" + +RegFitted = TypedDict( + "RegFitted", + { + "x": NDArray[np.float64], + "y": NDArray[np.float64], + "y.hat": NDArray[np.float64], + "NNS.ID": NDArray[np.str_], + "gradient": NDArray[np.float64], + "residuals": NDArray[np.float64], + "standard.errors": NDArray[np.float64], + "conf.int.pos": NotRequired[NDArray[np.float64]], + "conf.int.neg": NotRequired[NDArray[np.float64]], + }, +) +"""Per-observation fit table (R's ``$Fitted.xy``). + +The ``conf.int.*`` columns are added only when ``confidence_interval`` is set. +""" + +class RegEquation(TypedDict): + Variable: NDArray[np.str_] + Coefficient: NDArray[np.float64] +"""Dimension-reduction synthetic-regressor weights (R's ``$equation``).""" + +class RegXStar(TypedDict): + x: NDArray[np.float64] +"""Synthetic dimension-reduction regressor (R's ``$x.star``).""" + +RegResult = TypedDict( + "RegResult", + { + "R2": float, + "SE": float, + "Prediction.Accuracy": float | None, + "equation": "RegEquation | None", + "x.star": "RegXStar | None", + "derivative": DerivativeTable, + "Point.est": NDArray[np.float64], + "pred.int": "RegPredInt | None", + "regression.points": RegPoints, + "Fitted.xy": RegFitted, + }, +) +"""Univariate ``nns_reg`` result. + +``equation`` and ``x.star`` are populated only on the dimension-reduction path; +``Prediction.Accuracy`` only in class mode; ``pred.int`` only when +``confidence_interval`` is set; ``Point.est`` is empty unless ``point_est`` is given. +""" + @dataclass(frozen=True) class FactorDesign: @@ -31,6 +111,64 @@ class FactorDesign: feature_names: tuple[str, ...] +@overload +def nns_reg( + x: NDArray[Any], + y: NDArray[Any], + *, + factor_2_dummy: bool = ..., + order: Order = ..., + dim_red_method: object | None = ..., + tau: object | None = ..., + type: str | None = ..., + point_est: NDArray[np.float64] | float | None = ..., + return_values: bool = ..., + plot: bool = ..., + plot_regions: bool = ..., + residual_plot: bool = ..., + confidence_interval: float | None = ..., + threshold: float = ..., + n_best: object | None = ..., + smooth: bool = ..., + noise_reduction: NoiseReduction = ..., + dist: str = ..., + ncores: int | None = ..., + point_only: bool = ..., + multivariate_call: Literal[True], + class_levels: list[object] | None = ..., + factor_levels: Sequence[object] | Sequence[Sequence[object] | None] | None = ..., +) -> RegPoints: ... + + +@overload +def nns_reg( + x: NDArray[Any], + y: NDArray[Any], + *, + factor_2_dummy: bool = ..., + order: Order = ..., + dim_red_method: object | None = ..., + tau: object | None = ..., + type: str | None = ..., + point_est: NDArray[np.float64] | float | None = ..., + return_values: bool = ..., + plot: bool = ..., + plot_regions: bool = ..., + residual_plot: bool = ..., + confidence_interval: float | None = ..., + threshold: float = ..., + n_best: object | None = ..., + smooth: bool = ..., + noise_reduction: NoiseReduction = ..., + dist: str = ..., + ncores: int | None = ..., + point_only: bool = ..., + multivariate_call: Literal[False] = ..., + class_levels: list[object] | None = ..., + factor_levels: Sequence[object] | Sequence[Sequence[object] | None] | None = ..., +) -> RegResult | MRegResult: ... + + def nns_reg( x: NDArray[Any], y: NDArray[Any], @@ -56,8 +194,14 @@ def nns_reg( multivariate_call: bool = False, class_levels: list[object] | None = None, factor_levels: Sequence[object] | Sequence[Sequence[object] | None] | None = None, -) -> dict[str, Any]: - """Univariate numeric port of R's NNS.reg.""" +) -> RegResult | MRegResult | RegPoints: + """Univariate numeric port of R's NNS.reg. + + Returns :class:`RegResult` for 1-D ``x``. A 2-D ``x`` dispatches to + :func:`nns.multivariate_regression.nns_m_reg` and returns its + :class:`~nns.multivariate_regression.MRegResult`. The internal + ``multivariate_call=True`` contract returns bare :class:`RegPoints`. + """ _warn_unsupported( return_values=return_values is not True, ncores=ncores is not None, @@ -110,7 +254,7 @@ def nns_reg( dispatch_n_best = n_best if type_value == "class" and dispatch_n_best is None: dispatch_n_best = 1 - result = nns_m_reg( + m_reg_result = nns_m_reg( np.asarray(x_for_dispatch, dtype=np.float64), y_matrix_values, factor_2_dummy=False, @@ -128,7 +272,7 @@ def nns_reg( plot=plot, residual_plot=residual_plot, ) - return result + return m_reg_result # tau, threshold, n_best, and dist only apply to the dim-red and # multivariate paths dispatched above, mirroring R's NNS.reg. @@ -166,7 +310,7 @@ def nns_reg( def _maybe_render_reg( - result: dict[str, Any], + result: Mapping[str, object], *, plot: bool, plot_regions: bool, @@ -237,10 +381,10 @@ def _nns_reg_univariate_core( confidence_interval: float | None, multivariate_call: bool, class_mode: bool, - equation: dict[str, NDArray[np.float64] | NDArray[np.str_]] | None, - x_star: dict[str, NDArray[np.float64]] | None, + equation: RegEquation | None, + x_star: RegXStar | None, smooth: bool = False, -) -> dict[str, Any]: +) -> RegResult | RegPoints: dependence = _regression_dependence(x_values, y_values) dep_order = _dep_reduced_order(dependence, order, y_values.size) @@ -621,7 +765,7 @@ def _nns_reg_dimred( multivariate_call: bool, class_levels: list[object] | None = None, factor_levels: Sequence[object] | Sequence[Sequence[object] | None] | None = None, -) -> dict[str, Any]: +) -> RegResult | RegPoints: del n_best if factor_2_dummy: x, point_est, variable_names = _expand_factor_predictors_with_names( @@ -677,7 +821,7 @@ def __init__( self, x_star: NDArray[np.float64], point_est: NDArray[np.float64] | None, - equation: dict[str, NDArray[np.float64] | NDArray[np.str_]], + equation: RegEquation, ) -> None: self.x_star = x_star self.point_est = point_est @@ -773,8 +917,8 @@ def _dimred_projection( ) if len(names) != x.shape[1]: raise ValueError("variable_names must match the number of x columns.") - equation = { - "Variable": np.asarray([*names, "DENOMINATOR"]), + equation: RegEquation = { + "Variable": np.asarray([*names, "DENOMINATOR"], dtype=np.str_), "Coefficient": np.concatenate((coef, np.array([denominator], dtype=np.float64))), } return _DimredProjection(x_star=x_star, point_est=point_star, equation=equation) @@ -982,26 +1126,23 @@ def _partition_for_regression( dep_order: int | Literal["max"], requested_order: Order, noise: NoiseReduction, -) -> dict[str, Any]: +) -> PartResult: if dependence == 1.0 or dep_order == "max": if requested_order is None or dep_order == "max": return _max_order_part_map(x, y) - return cast(dict[str, Any], nns_part(x, y, order=int(dep_order), obs_req=0)) - return cast( - dict[str, Any], - nns_part( - x, - y, - noise_reduction=noise, - order=int(dep_order), - type="XONLY", - obs_req=0, - min_obs_stop=True, - ), + return nns_part(x, y, order=int(dep_order), obs_req=0) + return nns_part( + x, + y, + noise_reduction=noise, + order=int(dep_order), + type="XONLY", + obs_req=0, + min_obs_stop=True, ) -def _max_order_part_map(x: NDArray[np.float64], y: NDArray[np.float64]) -> dict[str, Any]: +def _max_order_part_map(x: NDArray[np.float64], y: NDArray[np.float64]) -> PartResult: quadrants = np.full(x.size, "q", dtype=str) seed_map = nns_part(x, y, order=1, obs_req=0) return { @@ -1158,7 +1299,7 @@ def _coefficients( rp_y: NDArray[np.float64], x: NDArray[np.float64], y: NDArray[np.float64], -) -> dict[str, NDArray[np.float64]]: +) -> DerivativeTable: if rp_x.size > 1: rise = np.diff(rp_y) run = np.diff(rp_x) @@ -1197,7 +1338,7 @@ def _fitted_values( y: NDArray[np.float64], rp_x: NDArray[np.float64], rp_y: NDArray[np.float64], - coeff: dict[str, NDArray[np.float64]], + coeff: DerivativeTable, order: Order, ) -> NDArray[np.float64]: if (order is not None and _is_fcl(order)) or ( @@ -1215,7 +1356,7 @@ def _predict_points( y: NDArray[np.float64], rp_x: NDArray[np.float64], rp_y: NDArray[np.float64], - coeff: dict[str, NDArray[np.float64]], + coeff: DerivativeTable, ) -> NDArray[np.float64]: reg_idx = _find_interval(point_est, rp_x, rightmost_closed=True) coef_idx = _find_interval(point_est, coeff["X.Lower.Range"], rightmost_closed=True) @@ -1251,7 +1392,7 @@ def _extrapolate_points( point_est_y: NDArray[np.float64], x: NDArray[np.float64], y: NDArray[np.float64], - coeff: dict[str, NDArray[np.float64]], + coeff: DerivativeTable, ) -> NDArray[np.float64]: out = point_est_y.astype(np.float64).copy() if not np.any((point_est > np.max(x)) | (point_est < np.min(x))): @@ -1305,8 +1446,8 @@ def _fitted_table( y: NDArray[np.float64], estimate: NDArray[np.float64], nns_ids: NDArray[np.str_], - coeff: dict[str, NDArray[np.float64]], -) -> dict[str, NDArray[np.float64] | NDArray[np.str_]]: + coeff: DerivativeTable, +) -> RegFitted: y_hat = estimate.copy() if np.any(~np.isfinite(y_hat)): replacement = _gravity(y_hat[np.isfinite(y_hat)]) @@ -1331,20 +1472,20 @@ def _fitted_table( def _apply_univariate_intervals( - fitted: dict[str, NDArray[np.float64] | NDArray[np.str_]], + fitted: RegFitted, point_values: NDArray[np.float64] | None, *, confidence_interval: float | None, class_mode: bool = False, -) -> dict[str, NDArray[np.float64]] | None: +) -> RegPredInt | None: if confidence_interval is None: return None alpha = (1.0 - float(confidence_interval)) / 2.0 - y_hat = cast(NDArray[np.float64], fitted["y.hat"]) - y = cast(NDArray[np.float64], fitted["y"]) - residuals = cast(NDArray[np.float64], fitted["residuals"]) - gradient = cast(NDArray[np.float64], fitted["gradient"]) + y_hat = fitted["y.hat"] + y = fitted["y"] + residuals = fitted["residuals"] + gradient = fitted["gradient"] conf_pos = np.empty_like(y_hat) conf_neg = np.empty_like(y_hat) @@ -1364,8 +1505,8 @@ def _apply_univariate_intervals( if point_values is None: return None - order = np.argsort(cast(NDArray[np.float64], fitted["x"]), kind="mergesort") - sorted_x = cast(NDArray[np.float64], fitted["x"])[order] + order = np.argsort(fitted["x"], kind="mergesort") + sorted_x = fitted["x"][order] row_indices: list[int] = [] for point in point_values: close = np.flatnonzero(np.isclose(sorted_x, point, rtol=1e-12, atol=1e-12)) @@ -1381,12 +1522,15 @@ def _apply_univariate_intervals( "pred.int.pos": np.array([], dtype=np.float64), } selected = np.asarray(row_indices, dtype=np.int64) - pred_int = { + pred_int: RegPredInt = { "pred.int.neg": pred_neg[order][selected], "pred.int.pos": pred_pos[order][selected], } if class_mode: - return {key: _round_class_interval(values) for key, values in pred_int.items()} + return { + "pred.int.neg": _round_class_interval(pred_int["pred.int.neg"]), + "pred.int.pos": _round_class_interval(pred_int["pred.int.pos"]), + } return pred_int diff --git a/src/nns/seasonality.py b/src/nns/seasonality.py index 907da374..3caed3fb 100644 --- a/src/nns/seasonality.py +++ b/src/nns/seasonality.py @@ -2,12 +2,30 @@ import math from collections import OrderedDict -from typing import SupportsInt, cast +from typing import SupportsInt, TypedDict, cast import numpy as np from numpy.typing import NDArray -SeasonalityResult = dict[str, object] +SeasonalityTable = TypedDict( + "SeasonalityTable", + { + "Period": NDArray[np.int64], + "Coefficient.of.Variation": NDArray[np.float64], + "Variable.Coefficient.of.Variation": NDArray[np.float64], + }, +) +"""Per-period seasonality diagnostics ordered by strength (R's ``$all.periods``).""" + +SeasonalityResult = TypedDict( + "SeasonalityResult", + { + "all.periods": SeasonalityTable, + "best.period": int, + "periods": NDArray[np.int64], + }, +) +"""``nns_seas`` result mirroring R's NNS.seas output names.""" _CacheKey = tuple[bytes, tuple[int, ...], bool] _CACHE_MAX_SIZE = 32 _CACHE: OrderedDict[_CacheKey, SeasonalityResult] = OrderedDict() @@ -374,15 +392,17 @@ def _clone_result(result: SeasonalityResult) -> SeasonalityResult: table = result["all.periods"] if not isinstance(table, dict): raise TypeError("Invalid seasonality result cache payload.") - cloned_table = { - "Period": np.asarray(table["Period"]).copy(), - "Coefficient.of.Variation": np.asarray(table["Coefficient.of.Variation"]).copy(), + cloned_table: SeasonalityTable = { + "Period": np.asarray(table["Period"], dtype=np.int64).copy(), + "Coefficient.of.Variation": np.asarray( + table["Coefficient.of.Variation"], dtype=np.float64 + ).copy(), "Variable.Coefficient.of.Variation": np.asarray( - table["Variable.Coefficient.of.Variation"] + table["Variable.Coefficient.of.Variation"], dtype=np.float64 ).copy(), } return { "all.periods": cloned_table, "best.period": int(cast(SupportsInt, result["best.period"])), - "periods": np.asarray(result["periods"]).copy(), + "periods": np.asarray(result["periods"], dtype=np.int64).copy(), } diff --git a/src/nns/stack.py b/src/nns/stack.py index 702de542..2f843eb6 100644 --- a/src/nns/stack.py +++ b/src/nns/stack.py @@ -2,7 +2,7 @@ import math from collections.abc import Callable, Sequence -from typing import Any, Literal, cast +from typing import Any, Literal, TypedDict, cast import numpy as np from numpy.typing import NDArray @@ -13,6 +13,7 @@ from nns.dependence import _gravity from nns.regression import ( Order, + RegResult, _expand_factor_predictors, _normalize_type, _prepare_y_values, @@ -22,8 +23,32 @@ ) Objective = Literal["min", "max"] + +StackPredInt = dict[str, NDArray[np.float64]] +"""Interval arrays keyed like R's stack pred.int output.""" Method = int | Sequence[int] -StackResult = dict[str, Any] +StackResult = TypedDict( + "StackResult", + { + "OBJfn.reg": float, + "NNS.reg.n.best": float, + "probability.threshold": float, + "OBJfn.dim.red": float, + "NNS.dim.red.threshold": float, + "reg": NDArray[np.float64], + "reg.pred.int": "StackPredInt | None", + "dim.red": NDArray[np.float64], + "dim.red.pred.int": "StackPredInt | None", + "stack": NDArray[np.float64], + "pred.int": "StackPredInt | None", + }, +) +"""``nns_stack`` result mirroring R's NNS.stack output names. + +``reg``/``dim.red`` carry each method's out-of-sample estimates and ``stack`` +the weighted combination; the ``*.pred.int`` entries are ``None`` unless +``pred_int`` is requested. +""" def nns_stack( @@ -299,7 +324,7 @@ def _evaluate_method2( fold_scores.append(float(scores[best_index])) if stack and methods == (1, 2): - fit = nns_reg( + fit = cast("RegResult", nns_reg( cv_x_train, cv_y_train, point_est=cv_x_test, @@ -308,7 +333,7 @@ def _evaluate_method2( order=order, dist=dist, point_only=False, - ) + )) train_star = cast(dict[str, NDArray[np.float64]], fit["x.star"])["x"] test_star = _xstar_for_points( fit, @@ -322,7 +347,7 @@ def _evaluate_method2( final_class_threshold = ( _threshold_mode(threshold_results) if type_value == "class" else math.nan ) - final_fit = nns_reg( + final_fit = cast("RegResult", nns_reg( x_train, y_train, point_est=x_test, @@ -333,7 +358,7 @@ def _evaluate_method2( point_only=False, confidence_interval=pred_int, type=type_value, - ) + )) fitted = cast(dict[str, NDArray[np.float64]], final_fit["Fitted.xy"]) fitted_yhat = fitted["y.hat"] prediction = _as_prediction(final_fit["Point.est"], x_test.shape[0]) @@ -579,7 +604,7 @@ def _fold_xstar( predicted = _class_threshold_round(predicted, threshold, cv_y_train) scores[idx] = objective_fn(predicted, cv_y_test) best_index = int(np.nanargmin(scores) if objective == "min" else np.nanargmax(scores)) - fit = nns_reg( + fit = cast("RegResult", nns_reg( cv_x_train, cv_y_train, point_est=cv_x_test, @@ -588,7 +613,7 @@ def _fold_xstar( order=order, dist=dist, point_only=False, - ) + )) return cast(dict[str, NDArray[np.float64]], fit["x.star"])["x"], _xstar_for_points( fit, cv_x_train, @@ -662,14 +687,14 @@ def _threshold_grid( elif isinstance(dim_red_method, str) and dim_red_method.lower() == "equal": return np.array([0.0], dtype=np.float64) else: - fit = nns_reg( + fit = cast("RegResult", nns_reg( x, y, dim_red_method=dim_red_method, order=order, dist=dist, point_only=True, - ) + )) equation = cast(dict[str, NDArray[np.float64]], fit["equation"]) scores = np.abs(np.round(equation["Coefficient"][:-1], 2)) scores = np.asarray(scores, dtype=np.float64) @@ -711,7 +736,7 @@ def _reg_point_est( def _xstar_for_points( - fit: dict[str, Any], + fit: RegResult, train_x: NDArray[np.float64], test_x: NDArray[np.float64], *, diff --git a/src/nns/stochastic_superiority.py b/src/nns/stochastic_superiority.py index 752641fd..bdeccb01 100644 --- a/src/nns/stochastic_superiority.py +++ b/src/nns/stochastic_superiority.py @@ -1,12 +1,10 @@ from __future__ import annotations -from typing import Any - import numpy as np from numpy.typing import NDArray from nns._native import native_fn -from nns.meboot import nns_meboot +from nns.meboot import MebootResult, nns_meboot from nns.var import lpm_var, upm_var @@ -99,7 +97,7 @@ def _omit_nan_numeric(x: NDArray[np.float64]) -> NDArray[np.float64]: return np.asarray(values[~np.isnan(values)], dtype=np.float64) -def _replicate_matrix(result: dict[str, Any]) -> NDArray[np.float64]: +def _replicate_matrix(result: MebootResult) -> NDArray[np.float64]: replicates = np.asarray(result.get("replicates"), dtype=np.float64) if replicates.ndim != 2: raise ValueError("NNS.meboot result does not contain a replicate matrix.") diff --git a/src/nns/var.py b/src/nns/var.py index 55c0cb77..aa8ddf21 100644 --- a/src/nns/var.py +++ b/src/nns/var.py @@ -3,7 +3,7 @@ import math from collections.abc import Sequence from numbers import Integral -from typing import Any, cast +from typing import Any, NotRequired, TypedDict, cast import numpy as np from numpy.typing import NDArray @@ -13,6 +13,16 @@ _R_OPTIMIZE_TOL = float(np.finfo(float).eps ** 0.25) +class VarResult(TypedDict): + """``nns_var`` result. ``h=0`` returns only the interpolated matrix and names.""" + + interpolated_and_extrapolated: NDArray[np.float64] + names: list[str] + relevant_variables: NotRequired[NDArray[Any]] + univariate: NotRequired[NDArray[np.float64]] + multivariate: NotRequired[NDArray[np.float64]] + ensemble: NotRequired[NDArray[np.float64]] + def nns_var( variables: NDArray[np.float64], @@ -26,7 +36,7 @@ def nns_var( status: bool = True, ncores: int | None = None, nowcast: bool = False, -) -> dict[str, Any]: +) -> VarResult: """Nonparametric VAR forecast for numeric matrix-like inputs. The public Python path returns plain arrays keyed like R's ``NNS.VAR`` output. @@ -101,7 +111,7 @@ def _var_interpolate_and_extrapolate( h: int, tau: int | Sequence[int] | Sequence[Sequence[int]] = 1, names: Sequence[str] | None = None, -) -> dict[str, object]: +) -> VarResult: """Interpolate missing values and generate univariate ARMA forecasts per column.""" vars_matrix = np.asarray(variables, dtype=np.float64) @@ -158,11 +168,12 @@ def _var_interpolate_and_extrapolate( plot=False, point_only=True, )["Point.est"] - if fitted_missing.size: + if fitted_missing is not None and fitted_missing.size: variable_interpolation[missing] = np.asarray(fitted_missing, dtype=np.float64) if h > 0: tau_i = _var_tau_for_variable(tau, j) + periods: NDArray[np.int64] | None try: periods = nns_seas( variable_interpolation, diff --git a/tests/invariants/test_boost.py b/tests/invariants/test_boost.py index 6b0c5d59..0df1ce28 100644 --- a/tests/invariants/test_boost.py +++ b/tests/invariants/test_boost.py @@ -7,6 +7,7 @@ import nns.boost as boost_module from nns import nns_boost +from nns.boost import BoostResult def test_nns_boost_shapes_and_feature_weights() -> None: @@ -322,7 +323,7 @@ def test_nns_boost_balance_retries_ordinary_fit_error(monkeypatch: pytest.Monkey original = boost_module._nns_boost_core calls = {"count": 0} - def fail_first(*args: Any, **kwargs: Any) -> dict[str, object]: + def fail_first(*args: Any, **kwargs: Any) -> BoostResult: calls["count"] += 1 if calls["count"] == 1: raise RuntimeError("ordinary fit failure") diff --git a/tests/parity/test_var.py b/tests/parity/test_var.py index a46a7173..a0a9de86 100644 --- a/tests/parity/test_var.py +++ b/tests/parity/test_var.py @@ -356,7 +356,7 @@ def test_public_nns_var_cor_matches_r( assert np.all(np.isfinite(actual_values)) _assert_public_numeric_close(actual_values, expected_values) assert np.array_equal( - cast(np.ndarray, actual_result["relevant_variables"]), + actual_result["relevant_variables"], cast(np.ndarray, expected_result["relevant_variables"]), ) @@ -381,7 +381,7 @@ def test_public_nns_var_cor_handles_missing_values_like_r() -> None: abs_tol=1e-8, ) assert np.array_equal( - cast(np.ndarray, actual_result["relevant_variables"]), + actual_result["relevant_variables"], cast(np.ndarray, expected_result["relevant_variables"]), ) @@ -413,7 +413,7 @@ def test_public_nns_var_nns_dep_matches_r() -> None: assert np.all(np.isfinite(actual_values)) _assert_public_numeric_close(actual_values, expected_values) assert np.array_equal( - cast(np.ndarray, actual_result["relevant_variables"]), + actual_result["relevant_variables"], cast(np.ndarray, expected_result["relevant_variables"]), ) @@ -445,7 +445,7 @@ def test_public_nns_var_nns_caus_matches_r() -> None: assert np.all(np.isfinite(actual_values)) _assert_public_numeric_close(actual_values, expected_values, rel_pct=1.0) assert np.array_equal( - cast(np.ndarray, actual_result["relevant_variables"]), + actual_result["relevant_variables"], cast(np.ndarray, expected_result["relevant_variables"]), ) @@ -477,7 +477,7 @@ def test_public_nns_var_all_matches_r() -> None: assert np.all(np.isfinite(actual_values)) _assert_public_numeric_close(actual_values, expected_values, rel_pct=1.0) assert np.array_equal( - cast(np.ndarray, actual_result["relevant_variables"]), + actual_result["relevant_variables"], cast(np.ndarray, expected_result["relevant_variables"]), ) diff --git a/uv.lock b/uv.lock index 5b45cd50..6248c8c8 100644 --- a/uv.lock +++ b/uv.lock @@ -1044,7 +1044,7 @@ wheels = [ [[package]] name = "ovvo-nns" -version = "1.0.9" +version = "1.1.0" source = { editable = "." } dependencies = [ { name = "matplotlib" },