Type public results as TypedDicts; bump to 1.1.0

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 |

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"

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"),

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)

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":

src/nns/meboot.py

@@ -1,14 +1,33 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import NotRequired, TypedDict
 
 import numpy as np
 from numpy.typing import NDArray
 
 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)

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