diff --git a/.mypyignore-todo b/.mypyignore-todo index c787d7ab..c9f2b790 100644 --- a/.mypyignore-todo +++ b/.mypyignore-todo @@ -1,4 +1,4 @@ scipy\.stats\.__all__ scipy\.stats\.(Normal|Uniform) -scipy\.stats\.(_distribution_infrastructure\.)?(abs|exp|log|make_distribution|order_statistic|truncate) -scipy\.stats\._distribution_infrastructure\.((Monotonic)?Transformed|Folded|Truncated|OrderStatistic)Distribution\.__init__ +scipy\.stats\.(_distribution_infrastructure\.)?(abs|exp|log|make_distribution|order_statistic) +scipy\.stats\._distribution_infrastructure\.(MonotonicTransformed|Folded|OrderStatistic)Distribution\.__init__ diff --git a/scipy-stubs/stats/__init__.pyi b/scipy-stubs/stats/__init__.pyi index 78046c7d..d9d39a53 100644 --- a/scipy-stubs/stats/__init__.pyi +++ b/scipy-stubs/stats/__init__.pyi @@ -7,7 +7,7 @@ from ._correlation import chatterjeexi from ._covariance import Covariance # TODO(jorenham) -from ._distribution_infrastructure import Mixture # , abs, exp, log, make_distribution, order_statistic, truncate +from ._distribution_infrastructure import Mixture, truncate # , abs, exp, log, make_distribution, order_statistic from ._entropy import differential_entropy, entropy from ._fit import fit, goodness_of_fit from ._hypotests import ( @@ -574,7 +574,7 @@ __all__ = [ "trim1", "trim_mean", "trimboth", - # "truncate", + "truncate", "truncexpon", "truncnorm", "truncpareto", diff --git a/scipy-stubs/stats/_distribution_infrastructure.pyi b/scipy-stubs/stats/_distribution_infrastructure.pyi index bac18902..013f4cd1 100644 --- a/scipy-stubs/stats/_distribution_infrastructure.pyi +++ b/scipy-stubs/stats/_distribution_infrastructure.pyi @@ -4,7 +4,7 @@ import abc from collections.abc import Mapping, Sequence, Set as AbstractSet from typing import Any, ClassVar, Final, Generic, Literal as L, TypeAlias, overload -from typing_extensions import LiteralString, Self, TypeVar, override +from typing_extensions import LiteralString, Never, Self, TypeVar, override import numpy as np import optype as op @@ -16,11 +16,11 @@ from ._probability_distribution import _BaseDistribution # TODO: # `__all__ = ["Mixture", "abs", "exp", "log", "make_distribution", "order_statistic", "truncate"] +_FloatT = TypeVar("_FloatT", bound=np.floating[Any]) + _ValidationPolicy: TypeAlias = L["skip_all"] | None _CachePolicy: TypeAlias = L["no_cache"] | None -_FloatT = TypeVar("_FloatT", bound=np.floating[Any]) - ### # TODO(jorenham): Generic dtype @@ -159,13 +159,14 @@ class _Parameterization: ### -_XT_co = TypeVar("_XT_co", bound=np.number[Any], default=np.float64, covariant=True) -_InexactT_co = TypeVar("_InexactT_co", bound=np.inexact[Any], default=np.float64, covariant=True) +_XT = TypeVar("_XT", bound=np.inexact[Any]) +_XT_co = TypeVar("_XT_co", bound=np.inexact[Any], default=np.float64, covariant=True) +_ShapeT0 = TypeVar("_ShapeT0", bound=tuple[int, ...]) _ShapeT0_co = TypeVar("_ShapeT0_co", bound=tuple[int, ...], default=tuple[int, ...], covariant=True) _DistrT_co = TypeVar( "_DistrT_co", - bound=ContinuousDistribution[np.number[Any]], - default=ContinuousDistribution[np.number[Any]], + bound=ContinuousDistribution[np.inexact[Any]], + default=ContinuousDistribution[np.inexact[Any]], covariant=True, ) @@ -197,6 +198,7 @@ class ContinuousDistribution(_BaseDistribution[_XT_co, _ShapeT0_co], Generic[_XT def cache_policy(self, /) -> _CachePolicy: ... @cache_policy.setter def cache_policy(self, cache_policy: _CachePolicy, /) -> None: ... + # def __init__( self, @@ -208,15 +210,15 @@ class ContinuousDistribution(_BaseDistribution[_XT_co, _ShapeT0_co], Generic[_XT ) -> None: ... # - def __neg__(self, /) -> ShiftedScaledDistribution[Self, _ShapeT0_co]: ... - def __abs__(self, /) -> FoldedDistribution[Self, _ShapeT0_co]: ... + def __neg__(self, /) -> ShiftedScaledDistribution[Self, _XT_co, _ShapeT0_co]: ... + def __abs__(self, /) -> FoldedDistribution[Self, _XT_co, _ShapeT0_co]: ... # TODO(jorenham): Accept `onp.ToFloatND`? - def __add__(self, rshift: onp.ToFloat, /) -> ShiftedScaledDistribution[Self, _ShapeT0_co]: ... - def __sub__(self, lshift: onp.ToFloat, /) -> ShiftedScaledDistribution[Self, _ShapeT0_co]: ... - def __mul__(self, scale: onp.ToFloat, /) -> ShiftedScaledDistribution[Self, _ShapeT0_co]: ... - def __truediv__(self, iscale: onp.ToFloat, /) -> ShiftedScaledDistribution[Self, _ShapeT0_co]: ... - def __pow__(self, exp: onp.ToInt, /) -> MonotonicTransformedDistribution[Self, _ShapeT0_co]: ... + def __add__(self, rshift: onp.ToFloat, /) -> ShiftedScaledDistribution[Self, _XT_co, _ShapeT0_co]: ... + def __sub__(self, lshift: onp.ToFloat, /) -> ShiftedScaledDistribution[Self, _XT_co, _ShapeT0_co]: ... + def __mul__(self, scale: onp.ToFloat, /) -> ShiftedScaledDistribution[Self, _XT_co, _ShapeT0_co]: ... + def __truediv__(self, iscale: onp.ToFloat, /) -> ShiftedScaledDistribution[Self, _XT_co, _ShapeT0_co]: ... + def __pow__(self, exp: onp.ToInt, /) -> MonotonicTransformedDistribution[Self, _XT_co, _ShapeT0_co]: ... __radd__ = __add__ __rsub__ = __sub__ __rmul__ = __mul__ @@ -262,47 +264,174 @@ class ContinuousDistribution(_BaseDistribution[_XT_co, _ShapeT0_co], Generic[_XT # -class TransformedDistribution(ContinuousDistribution[np.float64, _ShapeT0_co], Generic[_DistrT_co, _ShapeT0_co]): - # TODO(jorenham) - ... +class TransformedDistribution(ContinuousDistribution[_XT_co, _ShapeT0_co], Generic[_DistrT_co, _XT_co, _ShapeT0_co]): + def __init__( + self: TransformedDistribution[ContinuousDistribution[_XT, _ShapeT0], _XT, _ShapeT0], # nice trick, eh? + X: _DistrT_co, + /, + *args: Never, + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... -class TruncatedDistribution(TransformedDistribution[_DistrT_co, _ShapeT0_co], Generic[_DistrT_co, _ShapeT0_co]): +class MonotonicTransformedDistribution( + TransformedDistribution[_DistrT_co, _XT_co, _ShapeT0_co], + Generic[_DistrT_co, _XT_co, _ShapeT0_co], +): # TODO(jorenham) ... -class ShiftedScaledDistribution(TransformedDistribution[_DistrT_co, _ShapeT0_co], Generic[_DistrT_co, _ShapeT0_co]): - # TODO(jorenham) - ... +_DistrT0f = TypeVar("_DistrT0f", bound=ContinuousDistribution[np.floating[Any], tuple[()]]) +_DistrT1f = TypeVar("_DistrT1f", bound=ContinuousDistribution[np.floating[Any], tuple[int]]) +_DistrT2f = TypeVar("_DistrT2f", bound=ContinuousDistribution[np.floating[Any], tuple[int, int]]) +_DistrT3f = TypeVar("_DistrT3f", bound=ContinuousDistribution[np.floating[Any], tuple[int, int, int]]) +_DistrTNf = TypeVar("_DistrTNf", bound=ContinuousDistribution[np.floating[Any], tuple[int, ...]]) + +# still waiting on the intersection type PEP... +@overload +def truncate( + X: _DistrT0f, + lb: onp.ToFloat = ..., + ub: onp.ToFloat = ..., +) -> TruncatedDistribution[_DistrT0f, np.floating[Any], tuple[()]]: ... +@overload +def truncate( + X: _DistrT1f, + lb: onp.ToFloat | onp.ToFloatStrict1D = ..., + ub: onp.ToFloat | onp.ToFloatStrict1D = ..., +) -> TruncatedDistribution[_DistrT1f, np.floating[Any], tuple[int]]: ... +@overload +def truncate( + X: _DistrT2f, + lb: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D = ..., + ub: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D = ..., +) -> TruncatedDistribution[_DistrT2f, np.floating[Any], tuple[int, int]]: ... +@overload +def truncate( + X: _DistrT3f, + lb: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D | onp.ToFloatStrict3D = ..., + ub: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D | onp.ToFloatStrict3D = ..., +) -> TruncatedDistribution[_DistrT3f, np.floating[Any], tuple[int, int, int]]: ... +@overload +def truncate( + X: _DistrTNf, + lb: onp.ToFloat | onp.ToFloatND = ..., + ub: onp.ToFloat | onp.ToFloatND = ..., +) -> TruncatedDistribution[_DistrTNf, np.floating[Any], tuple[int, ...]]: ... +@overload +def truncate( + X: ContinuousDistribution[_XT, _ShapeT0], + lb: onp.ToFloat = ..., + ub: onp.ToFloat = ..., +) -> TruncatedDistribution[ContinuousDistribution[_XT, _ShapeT0], _XT, _ShapeT0]: ... +@overload +def truncate( + X: ContinuousDistribution[_XT], + lb: onp.ToFloat | onp.ToFloatND = ..., + ub: onp.ToFloat | onp.ToFloatND = ..., +) -> TruncatedDistribution[ContinuousDistribution[_XT], _XT]: ... + +class TruncatedDistribution(TransformedDistribution[_DistrT_co, _XT_co, _ShapeT0_co], Generic[_DistrT_co, _XT_co, _ShapeT0_co]): + lb: _XT_co | onp.ArrayND[_XT_co, _ShapeT0_co] + ub: _XT_co | onp.ArrayND[_XT_co, _ShapeT0_co] + + @overload + def __init__( + self: TruncatedDistribution[_DistrT0f, np.floating[Any], tuple[()]], + X: _DistrT0f, + /, + *args: Never, + lb: onp.ToFloat = ..., + ub: onp.ToFloat = ..., + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... + @overload + def __init__( + self: TruncatedDistribution[_DistrT1f, np.floating[Any], tuple[int]], + X: _DistrT1f, + /, + *args: Never, + lb: onp.ToFloat | onp.ToFloatStrict1D = ..., + ub: onp.ToFloat | onp.ToFloatStrict1D = ..., + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... + @overload + def __init__( + self: TruncatedDistribution[_DistrT2f, np.floating[Any], tuple[int, int]], + X: _DistrT2f, + /, + *args: Never, + lb: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D = ..., + ub: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D = ..., + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... + @overload + def __init__( + self: TruncatedDistribution[_DistrT3f, np.floating[Any], tuple[int, int, int]], + X: _DistrT3f, + /, + *args: Never, + lb: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D | onp.ToFloatStrict3D = ..., + ub: onp.ToFloat | onp.ToFloatStrict1D | onp.ToFloatStrict2D | onp.ToFloatStrict3D = ..., + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... + @overload + def __init__( + self: TruncatedDistribution[_DistrTNf, np.floating[Any], tuple[int, ...]], + X: _DistrTNf, + /, + *args: Never, + lb: onp.ToFloat | onp.ToFloatND = ..., + ub: onp.ToFloat | onp.ToFloatND = ..., + tol: opt.Just[float] | _Null = ..., + validation_policy: _ValidationPolicy = None, + cache_policy: _CachePolicy = None, + ) -> None: ... -class OrderStatisticDistribution(TransformedDistribution[_DistrT_co, _ShapeT0_co], Generic[_DistrT_co, _ShapeT0_co]): +class FoldedDistribution( + TransformedDistribution[_DistrT_co, _XT_co, _ShapeT0_co], + Generic[_DistrT_co, _XT_co, _ShapeT0_co], +): # TODO(jorenham) ... -class MonotonicTransformedDistribution(TransformedDistribution[_DistrT_co, _ShapeT0_co], Generic[_DistrT_co, _ShapeT0_co]): +class ShiftedScaledDistribution( + TransformedDistribution[_DistrT_co, _XT_co, _ShapeT0_co], + Generic[_DistrT_co, _XT_co, _ShapeT0_co], +): # TODO(jorenham) ... -class FoldedDistribution(TransformedDistribution[_DistrT_co, _ShapeT0_co], Generic[_DistrT_co, _ShapeT0_co]): +class OrderStatisticDistribution(TransformedDistribution[_DistrT_co, np.float64, _ShapeT0_co], Generic[_DistrT_co, _ShapeT0_co]): # TODO(jorenham) ... -class Mixture(_BaseDistribution[_InexactT_co, tuple[()]], Generic[_InexactT_co]): +class Mixture(_BaseDistribution[_XT_co, tuple[()]], Generic[_XT_co]): _shape: tuple[()] - _dtype: np.dtype[_InexactT_co] - _components: Sequence[ContinuousDistribution[_InexactT_co, tuple[()]]] - _weights: onp.Array1D[_InexactT_co] + _dtype: np.dtype[_XT_co] + _components: Sequence[ContinuousDistribution[_XT_co, tuple[()]]] + _weights: onp.Array1D[_XT_co] validation_policy: None @property - def components(self, /) -> list[ContinuousDistribution[_InexactT_co, tuple[()]]]: ... + def components(self, /) -> list[ContinuousDistribution[_XT_co, tuple[()]]]: ... @property - def weights(self, /) -> onp.Array1D[_InexactT_co]: ... + def weights(self, /) -> onp.Array1D[_XT_co]: ... # def __init__( self, /, - components: Sequence[ContinuousDistribution[_InexactT_co, tuple[()]]], + components: Sequence[ContinuousDistribution[_XT_co, tuple[()]]], *, weights: onp.ToFloat1D | None = None, ) -> None: ...