Merge pull request #354 from jorenham/faster-tests

faster doctests by 20-30 secs

lmo/_poly.py

@@ -9,7 +9,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, TypeAlias, TypeVar, cast, overload
+from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, cast, overload
 
 import numpy as np
 import numpy.polynomial as npp
@@ -51,19 +51,24 @@ def __dir__() -> tuple[str, ...]:
 
 
 @overload
-def eval_sh_jacobi(n: int, a: onp.ToFloat, b: onp.ToFloat, x: onp.ToFloat) -> float: ...
+def eval_sh_jacobi(
+    n: int | lmt.Integer,
+    a: float,
+    b: float,
+    x: float | np.floating[Any],
+) -> float: ...
 @overload
 def eval_sh_jacobi(
-    n: int,
-    a: onp.ToFloat,
-    b: onp.ToFloat,
+    n: int | lmt.Integer,
+    a: float,
+    b: float,
     x: onp.Array[_T_shape, lmt.Floating],
 ) -> onp.Array[_T_shape, np.float64]: ...
-def eval_sh_jacobi(
+def eval_sh_jacobi(  # noqa: C901
     n: int | lmt.Integer,
-    a: onp.ToFloat,
-    b: onp.ToFloat,
-    x: onp.ToFloat | onp.Array[_T_shape, lmt.Floating],
+    a: float,
+    b: float,
+    x: float | np.floating[Any] | onp.Array[_T_shape, lmt.Floating],
 ) -> float | onp.Array[_T_shape, np.float64]:
     """
     Fast evaluation of the n-th shifted Jacobi polynomial.
@@ -73,27 +78,23 @@ def eval_sh_jacobi(
     if n == 0:
         return 1
 
-    x = np.asarray(x)[()]
-    u = 2 * x - 1
-
-    a = float(a)
-    b = float(b)
+    x = x if isinstance(x, np.ndarray) else float(x)
 
     if a == b == 0:
         if n == 1:
-            return u
+            return 2 * x - 1
+        if n > 4:
+            return sps.eval_sh_legendre(n, x)
 
         v = x * (x - 1)
 
         if n == 2:
             return 1 + 6 * v
         if n == 3:
-            return (1 + 10 * v) * u
+            return (1 + 10 * v) * (2 * x - 1)
         if n == 4:
             return 1 + 10 * v * (2 + 7 * v)
 
-        return sps.eval_sh_legendre(n, x)
-
     if n == 1:
         return (a + b + 2) * x - b - 1
     if n == 2:
@@ -110,7 +111,7 @@ def eval_sh_jacobi(
         ) / 6
 
     # don't use `eval_sh_jacobi`: https://github.com/scipy/scipy/issues/18988
-    return sps.eval_jacobi(n, a, b, u)
+    return sps.eval_jacobi(n, a, b, 2 * x - 1)
 
 
 def peaks_jacobi(n: int, a: float, b: float) -> onp.Array1D[np.float64]:

lmo/diagnostic.py

@@ -250,22 +250,23 @@ def l_moment_gof(
         >>> import lmo
         >>> import numpy as np
         >>> from lmo.diagnostic import l_moment_gof
-        >>> from scipy.stats import norm
+        >>> from scipy.special import ndtr as norm_cdf
         >>> rng = np.random.default_rng(12345)
-        >>> X = norm(13.12, 1.66)
         >>> n = 1_000
-        >>> x = X.rvs(n, random_state=rng)
+        >>> x = rng.standard_normal(n)
         >>> x_lm = lmo.l_moment(x, [1, 2, 3, 4])
-        >>> l_moment_gof(X, x_lm, n).pvalue
+        >>> x_lm.round(4)
+        array([ 0.0083,  0.573 , -0.0067,  0.0722])
+        >>> l_moment_gof(norm_cdf, x_lm, n).pvalue
         0.82597
 
-        Contaminated samples:
+        Contaminate 4% of the samples with Cauchy noise:
 
-        >>> y = 0.9 * x + 0.1 * rng.normal(X.mean(), X.std() * 10, n)
+        >>> y = np.r_[x[: n - n // 25], rng.standard_cauchy(n // 25)]
         >>> y_lm = lmo.l_moment(y, [1, 2, 3, 4])
-        >>> y_lm.round(3)
-        array([13.193, 1.286, 0.006, 0.168])
-        >>> l_moment_gof(X, y_lm, n).pvalue
+        >>> y_lm.round(4)
+        array([-0.0731,  0.6923, -0.0876,  0.1932])
+        >>> l_moment_gof(norm_cdf, y_lm, n).pvalue
         0.0
 
 
@@ -656,6 +657,7 @@ def l_ratio_bounds(
     return t_min.round(12)[()], t_max.round(12)[()]
 
 
+# TODO: faster doctest
 def rejection_point(
     influence_fn: Callable[[float], float],
     /,
@@ -691,16 +693,14 @@ def rejection_point(
         >>> rejection_point(if_l_loc_norm)
         nan
 
-        For the TL-location of the Gaussian distribution, and even for the
-        Student's t distribution with 4 degrees of freedom (3 also works, but
-        is very slow), they exist.
+        For the TL-location of the Gaussian and Logistic distributions, they exist.
 
         >>> influence_norm = dists.norm.l_moment_influence(1, trim=1)
-        >>> influence_t4 = dists.t(4).l_moment_influence(1, trim=1)
-        >>> influence_norm(np.inf), influence_t4(np.inf)
+        >>> influence_logi = dists.logistic.l_moment_influence(1, trim=1)
+        >>> influence_norm(np.inf), influence_logi(np.inf)
         (0.0, 0.0)
-        >>> rejection_point(influence_norm), rejection_point(influence_t4)
-        (6.0, 206.0)
+        >>> rejection_point(influence_norm), rejection_point(influence_logi)
+        (6.0, 21.0)
 
     Notes:
         Large rejection points (e.g. >1000) are unlikely to be found.

lmo/theoretical/_f_to_lcm.py

@@ -142,8 +142,8 @@ def l_comoment_from_pdf(
         into TL-moments by Elamir & Seheult (2003).
 
     Examples:
-        Find the L-coscale and TL-coscale matrices of the multivariate
-        Student's t distribution with 4 degrees of freedom:
+        Find the TL-coscale matrix of the multivariate Student's t distribution with 4
+        degrees of freedom:
 
         >>> from scipy.stats import multivariate_t
         >>> df = 4
@@ -152,26 +152,18 @@ def l_comoment_from_pdf(
         >>> X = multivariate_t(loc=loc, shape=cov, df=df)
 
         >>> from scipy.special import stdtr
-        >>> std = np.sqrt(np.diag(cov))
-        >>> cdf0 = lambda x: stdtr(df, (x - loc[0]) / std[0])
-        >>> cdf1 = lambda x: stdtr(df, (x - loc[1]) / std[1])
-
-        >>> l_cov = l_comoment_from_pdf(X.pdf, (cdf0, cdf1), 2)
-        >>> l_cov.round(4)
-        array([[1.0413, 0.3124],
-               [0.1562, 0.5207]])
+        >>> std0, std1 = np.sqrt(np.diag(cov))
+        >>> cdf0 = lambda x: stdtr(df, (x - loc[0]) / std0)
+        >>> cdf1 = lambda x: stdtr(df, (x - loc[1]) / std1)
+
         >>> tl_cov = l_comoment_from_pdf(X.pdf, (cdf0, cdf1), 2, trim=1)
         >>> tl_cov.round(4)
         array([[0.4893, 0.1468],
                [0.0734, 0.2447]])
 
         The (Pearson) correlation coefficient can be recovered in several ways:
 
-        >>> cov[0, 1] / np.sqrt(cov[0, 0] * cov[1, 1])  # "true" correlation
-        0.3
-        >>> l_cov[0, 1] / l_cov[0, 0]
-        0.3
-        >>> l_cov[1, 0] / l_cov[1, 1]
+        >>> cov[0, 1] / (std0 * std1)  # "true" correlation
         0.3
         >>> tl_cov[0, 1] / tl_cov[0, 0]
         0.3
@@ -221,7 +213,7 @@ def l_comoment_from_pdf(
     def integrand(*xs: float, i: int, j: int) -> float:
         q_j = cdfs[j](xs[j])
         p_j = eval_sh_jacobi(_r - 1, t, s, q_j)
-        x = np.asarray(xs, dtype=np.float64)
+        x = np.asarray(xs, dtype=float)
         return c * x[i] * q_j**s * (1 - q_j) ** t * p_j * pdf(x)
 
     from scipy.integrate import nquad