Skip to content

Commit 1656364

Browse files
jorenhamjorenham
authored
✨ stats: improve sem (#815)
2 parents 12924c3 + a63b10f commit 1656364

File tree

2 files changed

+154
-2
lines changed

2 files changed

+154
-2
lines changed

‎scipy-stubs/stats/_stats_py.pyi‎

Lines changed: 65 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -616,9 +616,72 @@ def relfreq(
616616
#
617617
def obrientransform(*samples: onp.ToFloatND) -> onp.Array2D[npc.floating] | onp.Array1D[np.object_]: ...
618618

619-
# TODO(jorenham): improve
619+
#
620+
@overload # 1d ~inexact64 | +integer, keepdims=False (default)
621+
def sem(
622+
a: onp.ToArrayStrict1D[complex, npc.inexact64 | npc.integer | np.bool_],
623+
axis: L[0, -1] | None = 0,
624+
ddof: int = 1,
625+
nan_policy: NanPolicy = "propagate",
626+
*,
627+
keepdims: L[False] = False,
628+
) -> np.float64: ...
629+
@overload # >1d ~inexact64 | +integer, axis: int (default)
630+
def sem(
631+
a: onp.CanArray[onp.AtLeast2D, np.dtype[npc.inexact64 | npc.integer | np.bool_]] | Sequence[onp.SequenceND[complex]],
632+
axis: int = 0,
633+
ddof: int = 1,
634+
nan_policy: NanPolicy = "propagate",
635+
*,
636+
keepdims: bool = False,
637+
) -> onp.ArrayND[np.float64]: ...
638+
@overload # ?d ~inexact64 | +integer, axis=None, keepdims=False (default)
639+
def sem(
640+
a: onp.ToArrayND[complex, npc.inexact64 | npc.integer | np.bool_],
641+
axis: None,
642+
ddof: int = 1,
643+
nan_policy: NanPolicy = "propagate",
644+
*,
645+
keepdims: L[False] = False,
646+
) -> np.float64: ...
647+
@overload # ?d ~inexact64 | +integer, keepdims=True
648+
def sem(
649+
a: onp.ToArrayND[complex, npc.inexact64 | npc.integer | np.bool_],
650+
axis: int | None = 0,
651+
ddof: int = 1,
652+
nan_policy: NanPolicy = "propagate",
653+
*,
654+
keepdims: L[True],
655+
) -> onp.ArrayND[np.float64]: ...
656+
@overload # 1d +complex, keepdims=False (default)
657+
def sem(
658+
a: onp.ToComplexStrict1D,
659+
axis: L[0, -1] | None = 0,
660+
ddof: int = 1,
661+
nan_policy: NanPolicy = "propagate",
662+
*,
663+
keepdims: L[False] = False,
664+
) -> npc.floating: ...
665+
@overload # >1d +complex, axis: int (default)
666+
def sem(
667+
a: onp.CanArray[onp.AtLeast2D, np.dtype[npc.number]],
668+
axis: int = 0,
669+
ddof: int = 1,
670+
nan_policy: NanPolicy = "propagate",
671+
*,
672+
keepdims: bool = False,
673+
) -> onp.ArrayND[npc.floating]: ...
674+
@overload # ?d +complex, axis=None, keepdims=False (default)
675+
def sem(
676+
a: onp.ToComplexND, axis: None, ddof: int = 1, nan_policy: NanPolicy = "propagate", *, keepdims: L[False] = False
677+
) -> npc.floating: ...
678+
@overload # ?d +complex, keepdims=True
679+
def sem(
680+
a: onp.ToComplexND, axis: int | None = 0, ddof: int = 1, nan_policy: NanPolicy = "propagate", *, keepdims: L[True]
681+
) -> onp.ArrayND[npc.floating]: ...
682+
@overload # ?d +complex
620683
def sem(
621-
a: onp.ToFloatND, axis: int | None = 0, ddof: int = 1, nan_policy: NanPolicy = "propagate", *, keepdims: bool = False
684+
a: onp.ToComplexND, axis: int | None = 0, ddof: int = 1, nan_policy: NanPolicy = "propagate", *, keepdims: bool = False
622685
) -> _FloatOrND: ...
623686

624687
# TODO(jorenham): improve

‎tests/stats/test_sem.pyi‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
# type-tests for `sem` from `stats/_stats_py.pyi`
2+
3+
from typing import assert_type
4+
5+
import numpy as np
6+
import optype.numpy as onp
7+
import optype.numpy.compat as npc
8+
9+
from scipy.stats import sem
10+
11+
py_i_1d: list[int]
12+
py_i_2d: list[list[int]]
13+
14+
py_f_1d: list[float]
15+
py_f_2d: list[list[float]]
16+
17+
bool_1d: onp.Array1D[np.bool_]
18+
bool_2d: onp.Array2D[np.bool_]
19+
20+
i16_1d: onp.Array1D[np.int16]
21+
i16_2d: onp.Array2D[np.int16]
22+
23+
f32_1d: onp.Array1D[np.float32]
24+
f32_2d: onp.Array2D[np.float32]
25+
26+
f64_1d: onp.Array1D[np.float64]
27+
f64_2d: onp.Array2D[np.float64]
28+
29+
c64_1d: onp.Array1D[np.complex64]
30+
c64_2d: onp.Array2D[np.complex64]
31+
32+
c128_1d: onp.Array1D[np.complex128]
33+
c128_2d: onp.Array2D[np.complex128]
34+
35+
###
36+
37+
assert_type(sem(py_i_1d), np.float64)
38+
assert_type(sem(py_f_1d), np.float64)
39+
assert_type(sem(bool_1d), np.float64)
40+
assert_type(sem(i16_1d), np.float64)
41+
assert_type(sem(f32_1d), npc.floating)
42+
assert_type(sem(f64_1d), np.float64)
43+
assert_type(sem(c64_1d), npc.floating)
44+
assert_type(sem(c128_1d), np.float64)
45+
46+
assert_type(sem(py_i_2d), onp.ArrayND[np.float64])
47+
assert_type(sem(py_f_2d), onp.ArrayND[np.float64])
48+
assert_type(sem(bool_2d), onp.ArrayND[np.float64])
49+
assert_type(sem(i16_2d), onp.ArrayND[np.float64])
50+
assert_type(sem(f32_2d), onp.ArrayND[npc.floating])
51+
assert_type(sem(f64_2d), onp.ArrayND[np.float64])
52+
assert_type(sem(c64_2d), onp.ArrayND[npc.floating])
53+
assert_type(sem(c128_2d), onp.ArrayND[np.float64])
54+
55+
assert_type(sem(py_i_1d, axis=None), np.float64)
56+
assert_type(sem(py_f_1d, axis=None), np.float64)
57+
assert_type(sem(bool_1d, axis=None), np.float64)
58+
assert_type(sem(i16_1d, axis=None), np.float64)
59+
assert_type(sem(f32_1d, axis=None), npc.floating)
60+
assert_type(sem(f64_1d, axis=None), np.float64)
61+
assert_type(sem(c64_1d, axis=None), npc.floating)
62+
assert_type(sem(c128_1d, axis=None), np.float64)
63+
64+
assert_type(sem(py_i_2d, axis=None), np.float64)
65+
assert_type(sem(py_f_2d, axis=None), np.float64)
66+
assert_type(sem(bool_2d, axis=None), np.float64)
67+
assert_type(sem(i16_2d, axis=None), np.float64)
68+
assert_type(sem(f32_2d, axis=None), npc.floating)
69+
assert_type(sem(f64_2d, axis=None), np.float64)
70+
assert_type(sem(c64_2d, axis=None), npc.floating)
71+
assert_type(sem(c128_2d, axis=None), np.float64)
72+
73+
assert_type(sem(py_i_1d, keepdims=True), onp.ArrayND[np.float64])
74+
assert_type(sem(py_f_1d, keepdims=True), onp.ArrayND[np.float64])
75+
assert_type(sem(bool_1d, keepdims=True), onp.ArrayND[np.float64])
76+
assert_type(sem(i16_1d, keepdims=True), onp.ArrayND[np.float64])
77+
assert_type(sem(f32_1d, keepdims=True), onp.ArrayND[npc.floating])
78+
assert_type(sem(f64_1d, keepdims=True), onp.ArrayND[np.float64])
79+
assert_type(sem(c64_1d, keepdims=True), onp.ArrayND[npc.floating])
80+
assert_type(sem(c128_1d, keepdims=True), onp.ArrayND[np.float64])
81+
82+
assert_type(sem(py_i_2d, keepdims=True), onp.ArrayND[np.float64])
83+
assert_type(sem(py_f_2d, keepdims=True), onp.ArrayND[np.float64])
84+
assert_type(sem(bool_2d, keepdims=True), onp.ArrayND[np.float64])
85+
assert_type(sem(i16_2d, keepdims=True), onp.ArrayND[np.float64])
86+
assert_type(sem(f32_2d, keepdims=True), onp.ArrayND[npc.floating])
87+
assert_type(sem(f64_2d, keepdims=True), onp.ArrayND[np.float64])
88+
assert_type(sem(c64_2d, keepdims=True), onp.ArrayND[npc.floating])
89+
assert_type(sem(c128_2d, keepdims=True), onp.ArrayND[np.float64])

0 commit comments

Comments
 (0)