✨ stats: improve zscore and gzscore

scipy-stubs/stats/_stats_py.pyi

@@ -89,6 +89,8 @@ __all__ = [
 
 _SCT = TypeVar("_SCT", bound=np.generic)
 
+_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...])
+_InexactT = TypeVar("_InexactT", bound=npc.inexact)
 _FloatT = TypeVar("_FloatT", bound=npc.floating, default=npc.floating)
 _RealT = TypeVar("_RealT", bound=_Real0D, default=_Real0D)
 _RealT_co = TypeVar("_RealT_co", bound=_Real0D, default=_Real0D, covariant=True)
@@ -684,17 +686,98 @@ def sem(
     a: onp.ToComplexND, axis: int | None = 0, ddof: int = 1, nan_policy: NanPolicy = "propagate", *, keepdims: bool = False
 ) -> _FloatOrND: ...
 
-# TODO(jorenham): improve
+# NOTE: keep in sync with `gzscore`
+@overload  # +integer, known shape
+def zscore(
+    a: onp.CanArray[_ShapeT, np.dtype[npc.integer | np.bool_]],
+    axis: int | None = 0,
+    ddof: int = 0,
+    nan_policy: NanPolicy = "propagate",
+) -> onp.ArrayND[np.float64, _ShapeT]: ...
+@overload  # known inexact dtype, known shape
+def zscore(
+    a: onp.CanArray[_ShapeT, np.dtype[_InexactT]], axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.ArrayND[_InexactT, _ShapeT]: ...
+@overload  # float 1d
+def zscore(
+    a: Sequence[float], axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array1D[np.float64]: ...
+@overload  # float 2d
+def zscore(
+    a: Sequence[Sequence[float]], axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array2D[np.float64]: ...
+@overload  # float 3d
+def zscore(
+    a: Sequence[Sequence[Sequence[float]]], axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array3D[np.float64]: ...
+@overload  # complex 1d
+def zscore(
+    a: Sequence[op.JustComplex], axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array1D[np.complex128]: ...
+@overload  # complex 2d
+def zscore(
+    a: Sequence[Sequence[op.JustComplex]], axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array2D[np.complex128]: ...
+@overload  # complex 3d
+def zscore(
+    a: Sequence[Sequence[Sequence[op.JustComplex]]], axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array3D[np.complex128]: ...
+@overload  # floating fallback
 def zscore(
     a: onp.ToFloatND, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
 ) -> onp.ArrayND[npc.floating]: ...
+@overload  # complex fallback
+def zscore(
+    a: onp.ToJustComplexND, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.ArrayND[npc.complexfloating]: ...
 
-# TODO(jorenham): improve
+# NOTE: keep in sync with `zscore`
+@overload  # +integer, known shape
+def gzscore(
+    a: onp.CanArray[_ShapeT, np.dtype[npc.integer | np.bool_]],
+    *,
+    axis: int | None = 0,
+    ddof: int = 0,
+    nan_policy: NanPolicy = "propagate",
+) -> onp.ArrayND[np.float64, _ShapeT]: ...
+@overload  # known inexact dtype, known shape
+def gzscore(
+    a: onp.CanArray[_ShapeT, np.dtype[_InexactT]], *, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.ArrayND[_InexactT, _ShapeT]: ...
+@overload  # float 1d
+def gzscore(
+    a: Sequence[float], *, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array1D[np.float64]: ...
+@overload  # float 2d
+def gzscore(
+    a: Sequence[Sequence[float]], *, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array2D[np.float64]: ...
+@overload  # float 3d
+def gzscore(
+    a: Sequence[Sequence[Sequence[float]]], *, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array3D[np.float64]: ...
+@overload  # complex 1d
+def gzscore(
+    a: Sequence[op.JustComplex], *, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array1D[np.complex128]: ...
+@overload  # complex 2d
+def gzscore(
+    a: Sequence[Sequence[op.JustComplex]], *, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array2D[np.complex128]: ...
+@overload  # complex 3d
+def gzscore(
+    a: Sequence[Sequence[Sequence[op.JustComplex]]], *, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.Array3D[np.complex128]: ...
+@overload  # floating fallback
 def gzscore(
     a: onp.ToFloatND, *, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
 ) -> onp.ArrayND[npc.floating]: ...
+@overload  # complex fallback
+def gzscore(
+    a: onp.ToJustComplexND, *, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
+) -> onp.ArrayND[npc.complexfloating]: ...
 
-# TODO(jorenham): improve
+# TODO(jorenham): improve like zscore
 @overload  # (real vector-like, real vector-like) -> floating vector
 def zmap(
     scores: onp.ToFloat1D, compare: onp.ToFloat1D, axis: int | None = 0, ddof: int = 0, nan_policy: NanPolicy = "propagate"
@@ -1568,7 +1651,7 @@ def lmoment(
     sample: onp.ToFloatStrict2D,
     order: _LMomentOrder,
     *,
-    axis: L[0, 1, -1, -2] = 0,
+    axis: int = 0,
     keepdims: onp.ToFalse = False,
     sorted: op.CanBool = False,
     standardize: op.CanBool = True,
@@ -1579,7 +1662,7 @@ def lmoment(
     sample: onp.ToFloatStrict2D,
     order: _LMomentOrder,
     *,
-    axis: L[0, 1, -1, -2] | None = 0,
+    axis: int | None = 0,
     keepdims: onp.ToTrue,
     sorted: op.CanBool = False,
     standardize: op.CanBool = True,
@@ -1590,7 +1673,7 @@ def lmoment(
     sample: onp.ToFloatStrict2D,
     order: _LMomentOrder1D | None = None,
     *,
-    axis: L[0, 1, -1, -2] = 0,
+    axis: int = 0,
     keepdims: onp.ToFalse = False,
     sorted: op.CanBool = False,
     standardize: op.CanBool = True,
@@ -1601,7 +1684,7 @@ def lmoment(
     sample: onp.ToFloatStrict2D,
     order: _LMomentOrder1D | None = None,
     *,
-    axis: L[0, 1, -1, -2] | None = 0,
+    axis: int | None = 0,
     keepdims: onp.ToTrue,
     sorted: op.CanBool = False,
     standardize: op.CanBool = True,
@@ -1612,7 +1695,7 @@ def lmoment(
     sample: onp.ToFloatStrict3D,
     order: _LMomentOrder,
     *,
-    axis: L[0, 1, 2, -1, -2, -3] = 0,
+    axis: int = 0,
     keepdims: onp.ToFalse = False,
     sorted: op.CanBool = False,
     standardize: op.CanBool = True,
@@ -1623,7 +1706,7 @@ def lmoment(
     sample: onp.ToFloatStrict3D,
     order: _LMomentOrder,
     *,
-    axis: L[0, 1, 2, -1, -2, -3] | None = 0,
+    axis: int | None = 0,
     keepdims: onp.ToTrue,
     sorted: op.CanBool = False,
     standardize: op.CanBool = True,
@@ -1634,7 +1717,7 @@ def lmoment(
     sample: onp.ToFloatStrict3D,
     order: _LMomentOrder1D | None = None,
     *,
-    axis: L[0, 1, 2, -1, -2, -3] = 0,
+    axis: int = 0,
     keepdims: onp.ToFalse = False,
     sorted: op.CanBool = False,
     standardize: op.CanBool = True,
@@ -1645,7 +1728,7 @@ def lmoment(
     sample: onp.ToFloatStrict3D,
     order: _LMomentOrder1D | None = None,
     *,
-    axis: L[0, 1, 2, -1, -2, -3] | None = 0,
+    axis: int | None = 0,
     keepdims: onp.ToTrue,
     sorted: op.CanBool = False,
     standardize: op.CanBool = True,

tests/stats/test_zscore.pyi

@@ -0,0 +1,70 @@
+# type-tests for `zscore` from `stats/_stats_py.pyi`
+
+from typing import assert_type
+
+import numpy as np
+import optype.numpy as onp
+
+from scipy.stats import zscore
+
+py_i_1d: list[int]
+py_i_2d: list[list[int]]
+
+py_f_1d: list[float]
+py_f_2d: list[list[float]]
+
+bool_1d: onp.Array1D[np.bool_]
+bool_2d: onp.Array2D[np.bool_]
+
+i16_1d: onp.Array1D[np.int16]
+i16_2d: onp.Array2D[np.int16]
+
+f32_1d: onp.Array1D[np.float32]
+f32_2d: onp.Array2D[np.float32]
+
+f64_1d: onp.Array1D[np.float64]
+f64_2d: onp.Array2D[np.float64]
+
+c64_1d: onp.Array1D[np.complex64]
+c64_2d: onp.Array2D[np.complex64]
+
+c128_1d: onp.Array1D[np.complex128]
+c128_2d: onp.Array2D[np.complex128]
+
+###
+
+assert_type(zscore(py_i_1d), onp.Array1D[np.float64])
+assert_type(zscore(py_f_1d), onp.Array1D[np.float64])
+assert_type(zscore(bool_1d), onp.Array1D[np.float64])
+assert_type(zscore(i16_1d), onp.Array1D[np.float64])
+assert_type(zscore(f32_1d), onp.Array1D[np.float32])
+assert_type(zscore(f64_1d), onp.Array1D[np.float64])
+assert_type(zscore(c64_1d), onp.Array1D[np.complex64])
+assert_type(zscore(c128_1d), onp.Array1D[np.complex128])
+
+assert_type(zscore(py_i_2d), onp.Array2D[np.float64])
+assert_type(zscore(py_f_2d), onp.Array2D[np.float64])
+assert_type(zscore(bool_2d), onp.Array2D[np.float64])
+assert_type(zscore(i16_2d), onp.Array2D[np.float64])
+assert_type(zscore(f32_2d), onp.Array2D[np.float32])
+assert_type(zscore(f64_2d), onp.Array2D[np.float64])
+assert_type(zscore(c64_2d), onp.Array2D[np.complex64])
+assert_type(zscore(c128_2d), onp.Array2D[np.complex128])
+
+assert_type(zscore(py_i_1d, axis=None), onp.Array1D[np.float64])
+assert_type(zscore(py_f_1d, axis=None), onp.Array1D[np.float64])
+assert_type(zscore(bool_1d, axis=None), onp.Array1D[np.float64])
+assert_type(zscore(i16_1d, axis=None), onp.Array1D[np.float64])
+assert_type(zscore(f32_1d, axis=None), onp.Array1D[np.float32])
+assert_type(zscore(f64_1d, axis=None), onp.Array1D[np.float64])
+assert_type(zscore(c64_1d, axis=None), onp.Array1D[np.complex64])
+assert_type(zscore(c128_1d, axis=None), onp.Array1D[np.complex128])
+
+assert_type(zscore(py_i_2d, axis=None), onp.Array2D[np.float64])
+assert_type(zscore(py_f_2d, axis=None), onp.Array2D[np.float64])
+assert_type(zscore(bool_2d, axis=None), onp.Array2D[np.float64])
+assert_type(zscore(i16_2d, axis=None), onp.Array2D[np.float64])
+assert_type(zscore(f32_2d, axis=None), onp.Array2D[np.float32])
+assert_type(zscore(f64_2d, axis=None), onp.Array2D[np.float64])
+assert_type(zscore(c64_2d, axis=None), onp.Array2D[np.complex64])
+assert_type(zscore(c128_2d, axis=None), onp.Array2D[np.complex128])