diff --git a/.mypyignore-todo b/.mypyignore-todo index cf72c50..fe71ef6 100644 --- a/.mypyignore-todo +++ b/.mypyignore-todo @@ -1,4 +1,2 @@ scipy\.stats\.__all__ scipy\.stats\.(Normal|Uniform) -scipy\.stats\.(_distribution_infrastructure\.)?(make_distribution|order_statistic) -scipy\.stats\._distribution_infrastructure\.(Folded|OrderStatistic)Distribution\.__init__ diff --git a/scipy-stubs/stats/__init__.pyi b/scipy-stubs/stats/__init__.pyi index ef0b0e1..44388cd 100644 --- a/scipy-stubs/stats/__init__.pyi +++ b/scipy-stubs/stats/__init__.pyi @@ -5,9 +5,7 @@ from ._bws_test import bws_test from ._censored_data import CensoredData from ._correlation import chatterjeexi from ._covariance import Covariance - -# TODO(jorenham) -from ._distribution_infrastructure import Mixture, abs, exp, log, truncate # make_distribution, order_statistic +from ._distribution_infrastructure import Mixture, abs, exp, log, make_distribution, order_statistic, truncate from ._entropy import differential_entropy, entropy from ._fit import fit, goodness_of_fit from ._hypotests import ( @@ -470,7 +468,7 @@ __all__ = [ "logser", "loguniform", "lomax", - # "make_distribution", + "make_distribution", "mannwhitneyu", "matrix_normal", "maxwell", @@ -506,7 +504,7 @@ __all__ = [ "normaltest", "norminvgauss", "obrientransform", - # "order_statistic", + "order_statistic", "ortho_group", "page_trend_test", "pareto", diff --git a/scipy-stubs/stats/_distribution_infrastructure.pyi b/scipy-stubs/stats/_distribution_infrastructure.pyi index ad92e59..a97dcef 100644 --- a/scipy-stubs/stats/_distribution_infrastructure.pyi +++ b/scipy-stubs/stats/_distribution_infrastructure.pyi @@ -4,13 +4,14 @@ import abc from collections.abc import Callable, Mapping, Sequence, Set as AbstractSet from typing import Any, ClassVar, Final, Generic, Literal as L, Protocol, TypeAlias, overload, type_check_only -from typing_extensions import LiteralString, Never, Self, TypeVar, override +from typing_extensions import LiteralString, Never, Self, TypeIs, TypeVar, override import numpy as np import optype as op import optype.numpy as onp import optype.typing as opt from scipy._typing import AnyShape, ToRNG +from ._distn_infrastructure import rv_continuous from ._probability_distribution import _BaseDistribution # TODO: @@ -86,6 +87,8 @@ class _ParameterField(Protocol[_FloatingT_co, _ShapeT0_co]): _null: Final[_Null] = ... +def _isnull(x: object) -> TypeIs[_Null | None]: ... + # TODO(jorenham): Generic dtype and shape class _Domain(abc.ABC): # NOTE: This is a `ClassVar[dict[str, float]]` that's overridden as instance attribute in `_SimpleDomain`. @@ -100,11 +103,7 @@ class _Domain(abc.ABC): @abc.abstractmethod def draw(self, /, n: int) -> onp.ArrayND[_FloatingT]: ... @abc.abstractmethod - def get_numerical_endpoints( - self, - /, - x: _ParamValues, - ) -> tuple[onp.ArrayND[_Float], onp.ArrayND[_Float]]: ... + def get_numerical_endpoints(self, /, x: _ParamValues) -> tuple[onp.ArrayND[_Float], onp.ArrayND[_Float]]: ... # TODO(jorenham): Generic dtype class _SimpleDomain(_Domain, metaclass=abc.ABCMeta): @@ -189,9 +188,7 @@ class _Parameterization: def copy(self, /) -> Self: ... def matches(self, /, parameters: AbstractSet[str]) -> bool: ... def validation( - self, - /, - parameter_values: Mapping[str, _Parameter], + self, /, parameter_values: Mapping[str, _Parameter] ) -> tuple[onp.ArrayND[np.bool_], np.dtype[np.floating[Any]]]: ... def draw( self, @@ -211,6 +208,8 @@ class ContinuousDistribution(_BaseDistribution[_FloatingT_co, _ShapeT0_co], Gene _not_implemented: Final[str] _original_parameters: dict[str, _FloatingT_co | onp.ArrayND[_FloatingT_co, _ShapeT0_co]] + _variable: _Parameter + @property def tol(self, /) -> float | np.float64 | _Null | None: ... @tol.setter @@ -409,6 +408,8 @@ class TransformedDistribution( ContinuousDistribution[_FloatingT_co, _ShapeT0_co], Generic[_CDistT_co, _FloatingT_co, _ShapeT0_co], ): + _dist: _CDistT_co # readonly + def __init__( self: TransformedDistribution[ContinuousDistribution[_FloatingT, _ShapeT0], _FloatingT, _ShapeT0], # nice trick, eh? X: _CDistT_co, @@ -490,8 +491,8 @@ class FoldedDistribution( ) -> None: ... class TruncatedDistribution( - TransformedDistribution[_CDistT_co, _FloatingT_co, _ShapeT0_co], - Generic[_CDistT_co, _FloatingT_co, _ShapeT0_co], + TransformedDistribution[_CDistT_co, np.floating[Any], _ShapeT0_co], + Generic[_CDistT_co, _ShapeT0_co], ): _lb_domain: ClassVar[_RealDomain] = ... _lb_param: ClassVar[_RealParameter] = ... @@ -499,12 +500,12 @@ class TruncatedDistribution( _ub_domain: ClassVar[_RealDomain] = ... _ub_param: ClassVar[_RealParameter] = ... - lb: _ParameterField[_FloatingT_co, _ShapeT0_co] - ub: _ParameterField[_FloatingT_co, _ShapeT0_co] + lb: _ParameterField[np.floating[Any], _ShapeT0_co] + ub: _ParameterField[np.floating[Any], _ShapeT0_co] @overload def __init__( - self: TruncatedDistribution[_CDistT0, np.floating[Any], tuple[()]], + self: TruncatedDistribution[_CDistT0, tuple[()]], X: _CDistT0, /, *args: Never, @@ -516,7 +517,7 @@ class TruncatedDistribution( ) -> None: ... @overload def __init__( - self: TruncatedDistribution[_CDistT1, np.floating[Any], tuple[int]], + self: TruncatedDistribution[_CDistT1, tuple[int]], X: _CDistT1, /, *args: Never, @@ -528,7 +529,7 @@ class TruncatedDistribution( ) -> None: ... @overload def __init__( - self: TruncatedDistribution[_CDistT2, np.floating[Any], tuple[int, int]], + self: TruncatedDistribution[_CDistT2, tuple[int, int]], X: _CDistT2, /, *args: Never, @@ -540,7 +541,7 @@ class TruncatedDistribution( ) -> None: ... @overload def __init__( - self: TruncatedDistribution[_CDistT3, np.floating[Any], tuple[int, int, int]], + self: TruncatedDistribution[_CDistT3, tuple[int, int, int]], X: _CDistT3, /, *args: Never, @@ -552,7 +553,7 @@ class TruncatedDistribution( ) -> None: ... @overload def __init__( - self: TruncatedDistribution[_CDistT, np.floating[Any], tuple[int, ...]], + self: TruncatedDistribution[_CDistT, tuple[int, ...]], X: _CDistT, /, *args: Never, @@ -563,6 +564,76 @@ class TruncatedDistribution( cache_policy: _CachePolicy = None, ) -> None: ... +# always float64 or longdouble +class OrderStatisticDistribution(TransformedDistribution[_CDistT_co, _Float, _ShapeT0_co], Generic[_CDistT_co, _ShapeT0_co]): + # these should actually be integral; but the `_IntegerDomain` isn't finished yet + _r_domain: ClassVar[_RealDomain] = ... + _r_param: ClassVar[_RealParameter] = ... + + _n_domain: ClassVar[_RealDomain] = ... + _n_param: ClassVar[_RealParameter] = ... + + @overload + def __init__( + self: OrderStatisticDistribution[_CDistT0, tuple[()]], + dist: _CDistT0, + /, + *args: Never, + r: onp.ToJustInt, + n: onp.ToJustInt, + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... + @overload + def __init__( + self: OrderStatisticDistribution[_CDistT1, tuple[int]], + dist: _CDistT1, + /, + *args: Never, + r: onp.ToJustInt | onp.ToJustIntStrict1D, + n: onp.ToJustInt | onp.ToJustIntStrict1D, + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... + @overload + def __init__( + self: OrderStatisticDistribution[_CDistT2, tuple[int, int]], + dist: _CDistT2, + /, + *args: Never, + r: onp.ToJustInt | onp.ToJustIntStrict1D | onp.ToJustIntStrict2D, + n: onp.ToJustInt | onp.ToJustIntStrict1D | onp.ToJustIntStrict2D, + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... + @overload + def __init__( + self: OrderStatisticDistribution[_CDistT3, tuple[int, int, int]], + dist: _CDistT3, + /, + *args: Never, + r: onp.ToJustInt | onp.ToJustIntStrict1D | onp.ToJustIntStrict2D | onp.ToJustIntStrict3D, + n: onp.ToJustInt | onp.ToJustIntStrict1D | onp.ToJustIntStrict2D | onp.ToJustIntStrict3D, + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... + @overload + def __init__( + self: OrderStatisticDistribution[_CDistT, tuple[int, ...]], + X: _CDistT, + /, + *args: Never, + r: onp.ToJustInt | onp.ToJustIntND, + n: onp.ToJustInt | onp.ToJustIntND, + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... + # without HKT there's no reasonable way tot determine the floating scalar type class MonotonicTransformedDistribution( TransformedDistribution[_CDistT_co, np.floating[Any], _ShapeT0_co], @@ -591,10 +662,6 @@ class MonotonicTransformedDistribution( cache_policy: _CachePolicy = None, ) -> None: ... -class OrderStatisticDistribution(TransformedDistribution[_CDistT_co, np.float64, _ShapeT0_co], Generic[_CDistT_co, _ShapeT0_co]): - # TODO(jorenham) - ... - class Mixture(_BaseDistribution[_FloatingT_co, tuple[()]], Generic[_FloatingT_co]): _shape: tuple[()] _dtype: np.dtype[_FloatingT_co] @@ -606,10 +673,8 @@ class Mixture(_BaseDistribution[_FloatingT_co, tuple[()]], Generic[_FloatingT_co def components(self, /) -> list[_CDist0[_FloatingT_co]]: ... @property def weights(self, /) -> onp.Array1D[_FloatingT_co]: ... - # def __init__(self, /, components: Sequence[_CDist0[_FloatingT_co]], *, weights: onp.ToFloat1D | None = None) -> None: ... - # @override def kurtosis(self, /, *, method: _SMomentMethod | None = None) -> _Float: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] @@ -622,31 +687,73 @@ def truncate( X: _CDistT0, lb: onp.ToFloat = ..., ub: onp.ToFloat = ..., -) -> TruncatedDistribution[_CDistT0, np.floating[Any], tuple[()]]: ... +) -> TruncatedDistribution[_CDistT0, tuple[()]]: ... @overload def truncate( X: _CDistT1, lb: onp.ToFloat | onp.ToFloatStrict1D = ..., ub: onp.ToFloat | onp.ToFloatStrict1D = ..., -) -> TruncatedDistribution[_CDistT1, np.floating[Any], tuple[int]]: ... +) -> TruncatedDistribution[_CDistT1, tuple[int]]: ... @overload def truncate( X: _CDistT2, lb: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D = ..., ub: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D = ..., -) -> TruncatedDistribution[_CDistT2, np.floating[Any], tuple[int, int]]: ... +) -> TruncatedDistribution[_CDistT2, tuple[int, int]]: ... @overload def truncate( X: _CDistT3, lb: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D | onp.ToFloatStrict3D = ..., ub: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D | onp.ToFloatStrict3D = ..., -) -> TruncatedDistribution[_CDistT3, np.floating[Any], tuple[int, int, int]]: ... +) -> TruncatedDistribution[_CDistT3, tuple[int, int, int]]: ... @overload def truncate( X: _CDistT, lb: onp.ToFloat | onp.ToFloatND = ..., ub: onp.ToFloat | onp.ToFloatND = ..., -) -> TruncatedDistribution[_CDistT, np.floating[Any], tuple[int, ...]]: ... +) -> TruncatedDistribution[_CDistT, tuple[int, ...]]: ... + +# +@overload +def order_statistic( + X: _CDistT0, + /, + *, + r: onp.ToJustInt, + n: onp.ToJustInt, +) -> OrderStatisticDistribution[_CDistT0, tuple[()]]: ... +@overload +def order_statistic( + X: _CDistT1, + /, + *, + r: onp.ToJustInt | onp.ToJustIntStrict1D, + n: onp.ToJustInt | onp.ToJustIntStrict1D, +) -> OrderStatisticDistribution[_CDistT1, tuple[int]]: ... +@overload +def order_statistic( + X: _CDistT2, + /, + *, + r: onp.ToJustInt | onp.ToJustIntStrict1D | onp.ToJustIntStrict2D, + n: onp.ToJustInt | onp.ToJustIntStrict1D | onp.ToJustIntStrict2D, +) -> OrderStatisticDistribution[_CDistT2, tuple[int, int]]: ... +@overload +def order_statistic( + X: _CDistT3, + /, + *, + r: onp.ToJustInt | onp.ToJustIntStrict1D | onp.ToJustIntStrict2D | onp.ToJustIntStrict3D, + n: onp.ToJustInt | onp.ToJustIntStrict1D | onp.ToJustIntStrict2D | onp.ToJustIntStrict3D, +) -> OrderStatisticDistribution[_CDistT3, tuple[int, int, int]]: ... +@overload +def order_statistic( + X: _CDistT, + /, + *, + r: onp.ToJustInt | onp.ToJustIntND, + n: onp.ToJustInt | onp.ToJustIntND, +) -> OrderStatisticDistribution[_CDistT, tuple[int, ...]]: ... # @overload @@ -683,3 +790,10 @@ def log(X: _CDistT2, /) -> MonotonicTransformedDistribution[_CDistT2, tuple[int, def log(X: _CDistT3, /) -> MonotonicTransformedDistribution[_CDistT3, tuple[int, int, int]]: ... @overload def log(X: _CDistT, /) -> MonotonicTransformedDistribution[_CDistT, tuple[int, ...]]: ... + +# NOTE: These currently don't support >0-d parameters, and it looks like they always return float64, regardless of dtype +@type_check_only +class CustomDistribution(ContinuousDistribution[np.float64, tuple[()]]): + _dtype: np.dtype[np.floating[Any]] # ignored + +def make_distribution(dist: rv_continuous) -> type[CustomDistribution]: ... diff --git a/scipy-stubs/stats/_probability_distribution.pyi b/scipy-stubs/stats/_probability_distribution.pyi index 2b6a8a3..611cc21 100644 --- a/scipy-stubs/stats/_probability_distribution.pyi +++ b/scipy-stubs/stats/_probability_distribution.pyi @@ -62,25 +62,20 @@ class _ProbabilityDistribution(Generic[_XT_co], metaclass=abc.ABCMeta): def mode(self, /, *, method: _ModeMethod) -> _XT_co | onp.ArrayND[_XT_co]: ... @abc.abstractmethod def sample( - self, - /, - shape: int | tuple[int, ...], - *, - method: _SampleMethod, - rng: _ToQRNG, + self, /, shape: int | tuple[int, ...], *, method: _SampleMethod, rng: _ToQRNG ) -> _XT_co | onp.Array[Any, _XT_co]: ... # `Any` shape is needed on `numpy<2.1` # @abc.abstractmethod - def mean(self, /, *, method: _RMomentMethod) -> _Float | _FloatND: ... + def mean(self, /, *, method: _RMomentMethod) -> _XT_co | onp.ArrayND[_XT_co]: ... @abc.abstractmethod - def variance(self, /, *, method: _CMomentMethod) -> _Float | _FloatND: ... + def variance(self, /, *, method: _CMomentMethod) -> _XT_co | onp.ArrayND[_XT_co]: ... @abc.abstractmethod - def standard_deviation(self, /, *, method: _CMomentMethod) -> _Float | _FloatND: ... + def standard_deviation(self, /, *, method: _CMomentMethod) -> _XT_co | onp.ArrayND[_XT_co]: ... @abc.abstractmethod - def skewness(self, /, *, method: _SMomentMethod) -> _Float | _FloatND: ... + def skewness(self, /, *, method: _SMomentMethod) -> _XT_co | onp.ArrayND[_XT_co]: ... @abc.abstractmethod - def kurtosis(self, /, *, method: _SMomentMethod) -> _Float | _FloatND: ... + def kurtosis(self, /, *, method: _SMomentMethod) -> _XT_co | onp.ArrayND[_XT_co]: ... # @overload @@ -203,45 +198,37 @@ class _BaseDistribution(_ProbabilityDistribution[_XT_co], Generic[_XT_co, _Shape # @overload - def mean(self: _Self0, /, *, method: _RMomentMethod = None) -> _Float: ... + def mean(self: _Self0[_XT], /, *, method: _RMomentMethod = None) -> _XT: ... @overload - def mean(self: _Self[Any, _ShapeT], /, *, method: _RMomentMethod = None) -> _FloatND[_ShapeT]: ... + def mean(self: _Self[_XT, _ShapeT], /, *, method: _RMomentMethod = None) -> onp.ArrayND[_XT, _ShapeT]: ... # @overload - def variance(self: _Self0, /, *, method: _CMomentMethod = None) -> _Float: ... + def variance(self: _Self0[_XT], /, *, method: _CMomentMethod = None) -> _XT: ... @overload - def variance(self: _Self[Any, _ShapeT], /, *, method: _CMomentMethod = None) -> _FloatND[_ShapeT]: ... + def variance(self: _Self[_XT, _ShapeT], /, *, method: _CMomentMethod = None) -> onp.ArrayND[_XT, _ShapeT]: ... # @overload - def standard_deviation(self: _Self0, /, *, method: _CMomentMethod = None) -> _Float: ... + def standard_deviation(self: _Self0[_XT], /, *, method: _CMomentMethod = None) -> _XT: ... @overload - def standard_deviation(self: _Self[Any, _ShapeT], /, *, method: _CMomentMethod = None) -> _FloatND[_ShapeT]: ... + def standard_deviation(self: _Self[_XT, _ShapeT], /, *, method: _CMomentMethod = None) -> onp.ArrayND[_XT, _ShapeT]: ... # @overload - def skewness(self: _Self0, /, *, method: _SMomentMethod = None) -> _Float: ... + def skewness(self: _Self0[_XT], /, *, method: _SMomentMethod = None) -> _XT: ... @overload - def skewness(self: _Self[Any, _ShapeT], /, *, method: _SMomentMethod = None) -> _FloatND[_ShapeT]: ... + def skewness(self: _Self[_XT, _ShapeT], /, *, method: _SMomentMethod = None) -> onp.ArrayND[_XT, _ShapeT]: ... # @overload def kurtosis( - self: _Self0, - /, - *, - method: _SMomentMethod = None, - convention: _KurtosisConvention = "non-excess", - ) -> _Float: ... + self: _Self0[_XT], /, *, method: _SMomentMethod = None, convention: _KurtosisConvention = "non-excess" + ) -> _XT: ... @overload def kurtosis( - self: _Self[Any, _ShapeT], - /, - *, - method: _SMomentMethod = None, - convention: _KurtosisConvention = "non-excess", - ) -> _FloatND[_ShapeT]: ... + self: _Self[_XT, _ShapeT], /, *, method: _SMomentMethod = None, convention: _KurtosisConvention = "non-excess" + ) -> onp.ArrayND[_XT, _ShapeT]: ... # @overload @@ -256,48 +243,23 @@ class _BaseDistribution(_ProbabilityDistribution[_XT_co], Generic[_XT_co, _Shape def moment(self: _Self0, /, order: onp.ToInt = 1, *, kind: L["standardized"], method: _SMomentMethod = None) -> _Float: ... @overload def moment( - self: _Self[Any, _ShapeT], - /, - order: onp.ToInt = 1, - kind: L["raw"] = "raw", - *, - method: _RMomentMethod = None, + self: _Self[Any, _ShapeT], /, order: onp.ToInt = 1, kind: L["raw"] = "raw", *, method: _RMomentMethod = None ) -> _FloatND[_ShapeT]: ... @overload def moment( - self: _Self[Any, _ShapeT], - /, - order: onp.ToInt, - kind: L["central"], - *, - method: _CMomentMethod = None, + self: _Self[Any, _ShapeT], /, order: onp.ToInt, kind: L["central"], *, method: _CMomentMethod = None ) -> _FloatND[_ShapeT]: ... @overload def moment( - self: _Self[Any, _ShapeT], - /, - order: onp.ToInt = 1, - *, - kind: L["central"], - method: _CMomentMethod = None, + self: _Self[Any, _ShapeT], /, order: onp.ToInt = 1, *, kind: L["central"], method: _CMomentMethod = None ) -> _FloatND[_ShapeT]: ... @overload def moment( - self: _Self[Any, _ShapeT], - /, - order: onp.ToInt, - kind: L["standardized"], - *, - method: _SMomentMethod = None, + self: _Self[Any, _ShapeT], /, order: onp.ToInt, kind: L["standardized"], *, method: _SMomentMethod = None ) -> _FloatND[_ShapeT]: ... @overload def moment( # pyright: ignore[reportIncompatibleMethodOverride] # pyright false positive bug - self: _Self[Any, _ShapeT], - /, - order: onp.ToInt = 1, - *, - kind: L["standardized"], - method: _SMomentMethod = None, + self: _Self[Any, _ShapeT], /, order: onp.ToInt = 1, *, kind: L["standardized"], method: _SMomentMethod = None ) -> _FloatND[_ShapeT]: ... #