Skip to content

Commit

Permalink
👽️ stats: 1.15.0 new truncate and TruncatedDistribution
Browse files Browse the repository at this point in the history
  • Loading branch information
@jorenham
jorenham committed Dec 20, 2024
1 parent cc15c07 commit ed01cea
Show file tree
Hide file tree
Showing 3 changed files with 164 additions and 35 deletions.
4 changes: 2 additions & 2 deletions .mypyignore-todo
View file
Original file line number Diff line number Diff line change
@@ -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__
4 changes: 2 additions & 2 deletions scipy-stubs/stats/__init__.pyi
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -574,7 +574,7 @@ __all__ = [
"trim1",
"trim_mean",
"trimboth",
# "truncate",
"truncate",
"truncexpon",
"truncnorm",
"truncpareto",
Expand Down
191 changes: 160 additions & 31 deletions scipy-stubs/stats/_distribution_infrastructure.pyi
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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,
)

Expand Down Expand Up @@ -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,
Expand All @@ -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__
Expand Down Expand Up @@ -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: ...
Expand Down

0 comments on commit ed01cea

Please sign in to comment.